karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

fix conflict

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
yewentao256
2026-07-20 13:29:12 +00:00
parent d123d2b00a ae10e855ab
commit 83089b35d8
271 changed files with 7199 additions and 4277 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/hardware_tests/cpu.yaml
+5 -1
View File
@@ -18,6 +18,8 @@ steps:
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
- tests/kernels/mamba/test_cpu_short_conv.py
- tests/kernels/mamba/test_causal_conv1d.py
- tests/kernels/mamba/test_mamba_ssm.py
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
@@ -28,7 +30,9 @@ steps:
pytest -x -v -s tests/kernels/test_onednn.py
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
# Note: SDE can't be downloaded from CI host because of AWS WAF
# - label: CPU-Compatibility Tests
.buildkite/release-pipeline.yaml
+341 -344
View File
@@ -137,7 +137,7 @@ steps:
- 'mv artifacts/reassembled/wheel "artifacts/dist/$$wheel_name"'
- "aws sts get-caller-identity"
- "VLLM_WHEEL_PLATFORM=macos bash .buildkite/scripts/upload-nightly-wheels.sh"
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
plugins:
- aws-assume-role-with-web-identity#v1.6.0:
role-arn: arn:aws:iam::936637512419:role/vllm-release-macos-wheel-uploader
@@ -590,373 +590,370 @@ steps:
#
# =============================================================================
- block: "Unblock ROCm wheel/image prerequisites"
- group: "Build ROCm Wheel / Image "
key: "build-rocm-wheel-image"
depends_on: ~
key: block-build-rocm
if: build.env("NIGHTLY") != "1"
steps:
# ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
- label: ":rocm: Build ROCm Base Image & Wheels"
id: build-rocm-base-wheels
depends_on: ~
agents:
queue: cpu_queue_release
commands:
- |
set -euo pipefail
# ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
- label: ":rocm: Build ROCm Base Image & Wheels"
id: build-rocm-base-wheels
depends_on:
- step: block-build-rocm
allow_failure: true
agents:
queue: cpu_queue_release
commands:
- |
set -euo pipefail
# Generate cache key
CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"
# Generate cache key
CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"
echo "========================================"
echo "ROCm Base Build Configuration"
echo "========================================"
echo " CACHE_KEY: $${CACHE_KEY}"
echo " ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
echo "========================================"
# Login to ECR
aws ecr-public get-login-password --region us-east-1 | \
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
IMAGE_EXISTS=false
WHEELS_EXIST=false
# Check ECR for Docker image
echo "========================================"
echo "ROCm Base Build Configuration"
echo "========================================"
echo " CACHE_KEY: $${CACHE_KEY}"
echo " ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
echo "========================================"
# Login to ECR
aws ecr-public get-login-password --region us-east-1 | \
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
IMAGE_EXISTS=false
WHEELS_EXIST=false
# Check ECR for Docker image
if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
IMAGE_EXISTS=true
echo "ECR image cache HIT"
fi
# Check S3 for wheels
WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
WHEELS_EXIST=true
echo "S3 wheels cache HIT"
fi
if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
IMAGE_EXISTS=true
echo "ECR image cache HIT"
fi
# Check S3 for wheels
WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
WHEELS_EXIST=true
echo "S3 wheels cache HIT"
fi
# Scenario 1: Both cached (best case)
if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
echo ""
echo "FULL CACHE HIT - Reusing both image and wheels"
echo ""
# Scenario 1: Both cached (best case)
if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
echo ""
echo "FULL CACHE HIT - Reusing both image and wheels"
echo ""
# Download wheels
.buildkite/scripts/cache-rocm-base-wheels.sh download
# Save ECR tag for downstream jobs
buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
# Scenario 2: Full rebuild needed
else
echo ""
echo " CACHE MISS - Building from scratch..."
echo ""
# Build full base image and push to ECR
DOCKER_BUILDKIT=1 docker buildx build \
--file docker/Dockerfile.rocm_base \
--tag "$${ECR_CACHE_TAG}" \
--build-arg USE_SCCACHE=1 \
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
--build-arg SCCACHE_REGION_NAME=us-west-2 \
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
--push \
.
# Build wheel extraction stage
DOCKER_BUILDKIT=1 docker buildx build \
--file docker/Dockerfile.rocm_base \
--tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
--target debs_wheel_release \
--build-arg USE_SCCACHE=1 \
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
--build-arg SCCACHE_REGION_NAME=us-west-2 \
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
--load \
.
# Extract and upload wheels
mkdir -p artifacts/rocm-base-wheels
cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
docker rm $${cid}
.buildkite/scripts/cache-rocm-base-wheels.sh upload
# Download wheels
.buildkite/scripts/cache-rocm-base-wheels.sh download
# Save ECR tag for downstream jobs
buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
# Scenario 2: Full rebuild needed
else
echo ""
echo " CACHE MISS - Building from scratch..."
echo ""
# Build full base image and push to ECR
DOCKER_BUILDKIT=1 docker buildx build \
--file docker/Dockerfile.rocm_base \
--tag "$${ECR_CACHE_TAG}" \
--build-arg USE_SCCACHE=1 \
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
--build-arg SCCACHE_REGION_NAME=us-west-2 \
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
--push \
.
# Build wheel extraction stage
DOCKER_BUILDKIT=1 docker buildx build \
--file docker/Dockerfile.rocm_base \
--tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
--target debs_wheel_release \
--build-arg USE_SCCACHE=1 \
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
--build-arg SCCACHE_REGION_NAME=us-west-2 \
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
--load \
.
# Extract and upload wheels
mkdir -p artifacts/rocm-base-wheels
cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
docker rm $${cid}
.buildkite/scripts/cache-rocm-base-wheels.sh upload
# Cache base docker image to ECR
docker push "$${ECR_CACHE_TAG}"
buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
echo ""
echo " Build complete - Image and wheels cached"
fi
# Cache base docker image to ECR
docker push "$${ECR_CACHE_TAG}"
buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
echo ""
echo " Build complete - Image and wheels cached"
fi
artifact_paths:
- "artifacts/rocm-base-wheels/*.whl"
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
artifact_paths:
- "artifacts/rocm-base-wheels/*.whl"
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
# ROCm Job 2: Build vLLM ROCm Wheel
- label: ":python: Build vLLM ROCm Wheel - x86_64"
id: build-rocm-vllm-wheel
depends_on:
- step: build-rocm-base-wheels
allow_failure: false
agents:
queue: cpu_queue_release
timeout_in_minutes: 180
commands:
# Download artifacts and prepare Docker image
- |
set -euo pipefail
# ROCm Job 2: Build vLLM ROCm Wheel
- label: ":python: Build vLLM ROCm Wheel - x86_64"
id: build-rocm-vllm-wheel
depends_on:
- step: build-rocm-base-wheels
allow_failure: false
agents:
queue: cpu_queue_release
timeout_in_minutes: 180
commands:
# Download artifacts and prepare Docker image
- |
set -euo pipefail
# Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
# This fixes version detection when tags are moved/force-pushed
echo "Fetching latest tags from origin..."
git fetch --tags --force origin
# Log tag information for debugging version detection
echo "========================================"
echo "Git Tag Verification"
echo "========================================"
echo "Current HEAD: $(git rev-parse HEAD)"
echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
echo ""
echo "Recent tags (pointing to commits near HEAD):"
git tag -l --sort=-creatordate | head -5
echo "setuptools_scm version detection:"
pip install -q setuptools_scm 2>/dev/null || true
python3 -c "import setuptools_scm; print(' Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo " (setuptools_scm not available in this environment)"
echo "========================================"
# Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
# This fixes version detection when tags are moved/force-pushed
echo "Fetching latest tags from origin..."
git fetch --tags --force origin
# Log tag information for debugging version detection
echo "========================================"
echo "Git Tag Verification"
echo "========================================"
echo "Current HEAD: $(git rev-parse HEAD)"
echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
echo ""
echo "Recent tags (pointing to commits near HEAD):"
git tag -l --sort=-creatordate | head -5
echo "setuptools_scm version detection:"
pip install -q setuptools_scm 2>/dev/null || true
python3 -c "import setuptools_scm; print(' Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo " (setuptools_scm not available in this environment)"
echo "========================================"
# Download wheel artifacts from current build
echo "Downloading wheel artifacts from current build"
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
# Download wheel artifacts from current build
echo "Downloading wheel artifacts from current build"
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
# Get ECR image tag from metadata (set by build-rocm-base-wheels)
ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
if [ -z "$${ECR_IMAGE_TAG}" ]; then
echo "ERROR: rocm-base-image-tag metadata not found"
echo "This should have been set by the build-rocm-base-wheels job"
exit 1
fi
echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
# Login to ECR
aws ecr-public get-login-password --region us-east-1 | \
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
# Pull base Docker image from ECR
docker pull "$${ECR_IMAGE_TAG}"
echo "Loaded base image: $${ECR_IMAGE_TAG}"
# Prepare base wheels for Docker build context
mkdir -p docker/context/base-wheels
touch docker/context/base-wheels/.keep
cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
echo "Base wheels for vLLM build:"
ls -lh docker/context/base-wheels/
# Get ECR image tag from metadata (set by build-rocm-base-wheels)
ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
if [ -z "$${ECR_IMAGE_TAG}" ]; then
echo "ERROR: rocm-base-image-tag metadata not found"
echo "This should have been set by the build-rocm-base-wheels job"
exit 1
fi
echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
# Login to ECR
aws ecr-public get-login-password --region us-east-1 | \
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
# Pull base Docker image from ECR
docker pull "$${ECR_IMAGE_TAG}"
echo "Loaded base image: $${ECR_IMAGE_TAG}"
# Prepare base wheels for Docker build context
mkdir -p docker/context/base-wheels
touch docker/context/base-wheels/.keep
cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
echo "Base wheels for vLLM build:"
ls -lh docker/context/base-wheels/
echo "========================================"
echo "Building vLLM wheel with:"
echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
echo " BASE_IMAGE: $${ECR_IMAGE_TAG}"
echo "========================================"
echo "========================================"
echo "Building vLLM wheel with:"
echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
echo " BASE_IMAGE: $${ECR_IMAGE_TAG}"
echo "========================================"
# Build vLLM wheel using local checkout (REMOTE_VLLM=0)
DOCKER_BUILDKIT=1 docker build \
--file docker/Dockerfile.rocm \
--target export_vllm_wheel_release \
--output type=local,dest=rocm-dist \
--build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
--build-arg REMOTE_VLLM=0 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg USE_SCCACHE=1 \
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
--build-arg SCCACHE_REGION_NAME=us-west-2 \
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
.
echo "Built vLLM wheel:"
ls -lh rocm-dist/*.whl
# Copy wheel to artifacts directory
mkdir -p artifacts/rocm-vllm-wheel
cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
echo "Final vLLM wheel:"
ls -lh artifacts/rocm-vllm-wheel/
artifact_paths:
- "artifacts/rocm-vllm-wheel/*.whl"
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
# Build vLLM wheel using local checkout (REMOTE_VLLM=0)
DOCKER_BUILDKIT=1 docker build \
--file docker/Dockerfile.rocm \
--target export_vllm_wheel_release \
--output type=local,dest=rocm-dist \
--build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
--build-arg REMOTE_VLLM=0 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg USE_SCCACHE=1 \
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
--build-arg SCCACHE_REGION_NAME=us-west-2 \
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
.
echo "Built vLLM wheel:"
ls -lh rocm-dist/*.whl
# Copy wheel to artifacts directory
mkdir -p artifacts/rocm-vllm-wheel
cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
echo "Final vLLM wheel:"
ls -lh artifacts/rocm-vllm-wheel/
artifact_paths:
- "artifacts/rocm-vllm-wheel/*.whl"
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
# ROCm Job 3: Upload Wheels to S3
- label: ":s3: Upload ROCm Wheels to S3"
id: upload-rocm-wheels
depends_on:
- step: build-rocm-vllm-wheel
allow_failure: false
agents:
queue: cpu_queue_release
timeout_in_minutes: 60
commands:
# Download all wheel artifacts and run upload
- |
set -euo pipefail
# ROCm Job 3: Upload Wheels to S3
- label: ":s3: Upload ROCm Wheels to S3"
id: upload-rocm-wheels
depends_on:
- step: build-rocm-vllm-wheel
allow_failure: false
agents:
queue: cpu_queue_release
timeout_in_minutes: 60
commands:
# Download all wheel artifacts and run upload
- |
set -euo pipefail
# Download artifacts from current build
echo "Downloading artifacts from current build"
# buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
# buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
# Download artifacts from current build
echo "Downloading artifacts from current build"
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
# # Run upload script
bash .buildkite/scripts/upload-rocm-wheels.sh
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
# Run upload script
bash .buildkite/scripts/upload-rocm-wheels.sh
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
# ROCm Job 4: Annotate ROCm Wheel Release
- label: ":memo: Annotate ROCm wheel release"
id: annotate-rocm-release
depends_on:
- upload-rocm-wheels
agents:
queue: cpu_queue_release
commands:
- "bash .buildkite/scripts/annotate-rocm-release.sh"
env:
S3_BUCKET: "vllm-wheels"
# ROCm Job 4: Annotate ROCm Wheel Release
- label: ":memo: Annotate ROCm wheel release"
id: annotate-rocm-release
depends_on:
- upload-rocm-wheels
agents:
queue: cpu_queue_release
commands:
- "bash .buildkite/scripts/annotate-rocm-release.sh"
env:
S3_BUCKET: "vllm-wheels"
# ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
# This is the job to create https://wheels.vllm.ai/rocm/ index allowing
# users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
- block: "Generate Root Index for ROCm Wheels for Release"
key: block-generate-root-index-rocm-wheels
depends_on: upload-rocm-wheels
# ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
# This is the job to create https://wheels.vllm.ai/rocm/ index allowing
# users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
- block: "Generate Root Index for ROCm Wheels for Release"
key: block-generate-root-index-rocm-wheels
depends_on: upload-rocm-wheels
- label: ":package: Generate Root Index for ROCm Wheels for Release"
depends_on: block-generate-root-index-rocm-wheels
id: generate-root-index-rocm-wheels
agents:
queue: cpu_queue_release
commands:
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
env:
S3_BUCKET: "vllm-wheels"
VARIANT: "rocm723"
- label: ":package: Generate Root Index for ROCm Wheels for Release"
depends_on: block-generate-root-index-rocm-wheels
id: generate-root-index-rocm-wheels
agents:
queue: cpu_queue_release
commands:
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
env:
S3_BUCKET: "vllm-wheels"
VARIANT: "rocm723"
# ROCm Job 6: Build ROCm Release Docker Image
- label: ":docker: Build release image - x86_64 - ROCm"
id: build-rocm-release-image
depends_on:
- step: block-build-release-images
allow_failure: true
- step: build-rocm-base-wheels
allow_failure: false
agents:
queue: cpu_queue_release
timeout_in_minutes: 60
commands:
- |
set -euo pipefail
# Login to ECR
aws ecr-public get-login-password --region us-east-1 | \
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
# Get ECR image tag from metadata (set by build-rocm-base-wheels)
ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
if [ -z "$${ECR_IMAGE_TAG}" ]; then
echo "ERROR: rocm-base-image-tag metadata not found"
echo "This should have been set by the build-rocm-base-wheels job"
exit 1
fi
echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
# Pull base Docker image from ECR
docker pull "$${ECR_IMAGE_TAG}"
echo "Loaded base image: $${ECR_IMAGE_TAG}"
# Pass the base image ECR tag to downstream steps (nightly publish)
buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}"
echo "========================================"
echo "Building vLLM ROCm release image with:"
echo " BASE_IMAGE: $${ECR_IMAGE_TAG}"
echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
echo "========================================"
# Build vLLM ROCm release image using cached base
DOCKER_BUILDKIT=1 docker build \
--build-arg max_jobs=16 \
--build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
--build-arg USE_SCCACHE=1 \
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
--build-arg SCCACHE_REGION_NAME=us-west-2 \
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
--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
# ROCm Job 6: Build ROCm Release Docker Image
- label: ":docker: Build release image - x86_64 - ROCm"
id: build-rocm-release-image
depends_on:
- step: block-build-release-images
allow_failure: true
- step: build-rocm-base-wheels
allow_failure: false
agents:
queue: cpu_queue_release
timeout_in_minutes: 60
commands:
- |
set -euo pipefail
# Login to ECR
aws ecr-public get-login-password --region us-east-1 | \
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
# Get ECR image tag from metadata (set by build-rocm-base-wheels)
ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
if [ -z "$${ECR_IMAGE_TAG}" ]; then
echo "ERROR: rocm-base-image-tag metadata not found"
echo "This should have been set by the build-rocm-base-wheels job"
exit 1
fi
echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
# Pull base Docker image from ECR
docker pull "$${ECR_IMAGE_TAG}"
echo "Loaded base image: $${ECR_IMAGE_TAG}"
# Pass the base image ECR tag to downstream steps (nightly publish)
buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}"
echo "========================================"
echo "Building vLLM ROCm release image with:"
echo " BASE_IMAGE: $${ECR_IMAGE_TAG}"
echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
echo "========================================"
# Build vLLM ROCm release image using cached base
DOCKER_BUILDKIT=1 docker build \
--build-arg max_jobs=16 \
--build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
--build-arg USE_SCCACHE=1 \
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
--build-arg SCCACHE_REGION_NAME=us-west-2 \
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
--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 ""
echo " Successfully built and pushed ROCm release image"
echo " Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
echo ""
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
echo ""
echo " Successfully built and pushed ROCm release image"
echo " Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
echo ""
env:
DOCKER_BUILDKIT: "1"
S3_BUCKET: "vllm-wheels"
- label: "Publish nightly XPU image to DockerHub"
depends_on:
- create-manifest-xpu
if: build.env("NIGHTLY") == "1"
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
- "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
plugins:
- docker-login#v3.0.0:
username: vllmbot
password-env: DOCKERHUB_TOKEN
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- label: "Publish nightly XPU image to DockerHub"
depends_on:
- create-manifest-xpu
if: build.env("NIGHTLY") == "1"
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
- "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
plugins:
- docker-login#v3.0.0:
username: vllmbot
password-env: DOCKERHUB_TOKEN
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- label: "Publish nightly ROCm image to DockerHub"
depends_on:
- build-rocm-release-image
if: build.env("NIGHTLY") == "1"
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/push-nightly-builds-rocm.sh"
# Clean up old nightly builds (keep only last 14)
- "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm"
- "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm"
plugins:
- docker-login#v3.0.0:
username: vllmbot
password-env: DOCKERHUB_TOKEN
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- label: "Publish nightly ROCm image to DockerHub"
depends_on:
- build-rocm-release-image
if: build.env("NIGHTLY") == "1"
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/push-nightly-builds-rocm.sh"
# Clean up old nightly builds (keep only last 14)
- "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm"
- "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm"
plugins:
- docker-login#v3.0.0:
username: vllmbot
password-env: DOCKERHUB_TOKEN
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
# =============================================================================
# Publish to DockerHub and PyPI (at the end so all builds complete first)
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
+4 -1
View File
@@ -40,7 +40,9 @@ function cpu_tests() {
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
# skip tests requiring model downloads if HF_TOKEN is not set
# due to rate-limits
@@ -97,3 +99,4 @@ function cpu_tests() {
# All of CPU tests are expected to be finished less than 40 mins.
export -f cpu_tests
timeout 2h bash -c cpu_tests
.buildkite/test_areas/cuda.yaml
+2 -1
View File
@@ -18,6 +18,7 @@ steps:
- pytest -v -s cuda/test_platform_no_cuda_init.py
- label: Cudagraph
device: h200_35gb
key: cudagraph
timeout_in_minutes: 30
source_file_dependencies:
@@ -28,4 +29,4 @@ steps:
commands:
- pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
.buildkite/test_areas/entrypoints.yaml
+7
View File
@@ -3,6 +3,7 @@ depends_on:
- image-build
steps:
- label: Entrypoints Unit Tests
device: h200_35gb
key: entrypoints-unit-tests
timeout_in_minutes: 25
working_dir: "/vllm-workspace/tests"
@@ -15,6 +16,7 @@ steps:
- pytest -v -s entrypoints/weight_transfer
- label: Entrypoints Integration (LLM)
device: h200_35gb
key: entrypoints-integration-llm
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
@@ -55,6 +57,7 @@ steps:
- image-build-amd
- label: Entrypoints Integration (API Server OpenAI - Part 1)
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-1
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
@@ -73,6 +76,7 @@ steps:
- image-build-amd
- label: Entrypoints Integration (API Server OpenAI - Part 2)
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-2
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
@@ -92,6 +96,7 @@ steps:
- image-build-amd
- label: Entrypoints Integration (API Server Generate)
device: h200_35gb
key: entrypoints-integration-api-server-generate
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
@@ -114,6 +119,7 @@ steps:
- image-build-amd
- label: Entrypoints Integration (Responses API)
device: h200_35gb
key: entrypoints-integration-responses-api
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
@@ -148,6 +154,7 @@ steps:
- pytest -v -s entrypoints/multimodal
- label: Entrypoints Integration (Pooling)
device: h200_35gb
key: entrypoints-integration-pooling
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
.buildkite/test_areas/kernels.yaml
+9
View File
@@ -15,6 +15,7 @@ steps:
- pytest -v -s tests/kernels/ir
- label: Kernels Core Operation Test
device: h200_35gb
key: kernels-core-operation-test
timeout_in_minutes: 120
source_file_dependencies:
@@ -163,6 +164,7 @@ steps:
- image-build-amd
- label: Kernels Mamba Test
device: h200_35gb
key: kernels-mamba-test
timeout_in_minutes: 40
source_file_dependencies:
@@ -235,6 +237,11 @@ steps:
- vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py
- vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py
- vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py
- vllm/cute_utils/
- vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/
- vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py
- tests/kernels/mamba/test_gdn_prefill_cutedsl.py
- tests/kernels/test_bf16x3_router_gemm_cutedsl.py
- tests/kernels/test_ll_bf16_gemm.py
- tests/kernels/test_top_k_per_row.py
commands:
@@ -264,6 +271,8 @@ steps:
- pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
- pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py
- pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
- pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py
- pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py
- pytest -v -s tests/kernels/test_ll_bf16_gemm.py
# e2e
- pytest -v -s tests/models/quantization/test_nvfp4.py
.buildkite/test_areas/lm_eval.yaml
+22
View File
@@ -78,6 +78,28 @@ steps:
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt
- label: LM Eval PCP (4xB200)
key: lm-eval-pcp-4xb200
timeout_in_minutes: 360
device: b200-k8s
num_devices: 4
optional: true
source_file_dependencies:
- csrc/
- tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
- tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
- tests/evals/gsm8k/configs/models-pcp.txt
- vllm/model_executor/layers/quantization
- vllm/config/parallel.py
- vllm/distributed/parallel_state.py
- vllm/model_executor/layers/attention/mla_attention.py
- vllm/model_executor/layers/attention/pcp.py
- vllm/v1/worker/gpu/model_runner.py
- vllm/v1/worker/gpu/pcp_manager.py
autorun_on_main: true
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt
- label: LM Eval Large Models EP (2xB200)
key: lm-eval-large-models-ep-2xb200
timeout_in_minutes: 60
.buildkite/test_areas/misc.yaml
+2 -1
View File
@@ -64,8 +64,9 @@ steps:
- image-build-amd
- label: V1 Core + KV + Metrics
device: h200_35gb
key: v1-core-kv-metrics
timeout_in_minutes: 60
timeout_in_minutes: 80
source_file_dependencies:
- vllm/config/
- vllm/distributed/
.buildkite/test_areas/model_executor.yaml
+1
View File
@@ -3,6 +3,7 @@ depends_on:
- image-build
steps:
- label: Model Executor
device: h200_35gb
key: model-executor
timeout_in_minutes: 45
source_file_dependencies:
.buildkite/test_areas/models_language.yaml
+18 -3
View File
@@ -21,6 +21,7 @@ steps:
- image-build-amd
- label: Language Models Tests (Extra Standard) %N
device: h200_35gb
key: language-models-tests-extra-standard
timeout_in_minutes: 40
source_file_dependencies:
@@ -51,8 +52,8 @@ steps:
- tests/models/language/pooling/test_classification.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
- label: Language Models Tests (Hybrid) %N
device: h200_35gb
key: language-models-tests-hybrid
timeout_in_minutes: 65
source_file_dependencies:
@@ -63,8 +64,8 @@ steps:
# Note: also needed to run plamo2 model in vLLM
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
# Shard hybrid language model tests
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
# Shard the hybrid language model tests that are numerically stable on Hopper.
- pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
parallelism: 2
mirror:
amd:
@@ -77,6 +78,20 @@ steps:
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
# Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD
# autotuning (https://github.com/vllm-project/vllm/issues/25194). Keep this one
# correctness test on L4 until its H200 output matches the Transformers reference.
- label: Language Models Tests (Granite L4 Compatibility)
key: language-models-tests-granite-l4-compatibility
timeout_in_minutes: 65
source_file_dependencies:
- vllm/
- tests/models/language/generation
commands:
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
- pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
- label: Language Models Test (Extended Generation) # 80min
device: h200_35gb
key: language-models-test-extended-generation
.buildkite/test_areas/models_multimodal.yaml
+1
View File
@@ -119,6 +119,7 @@ steps:
- vllm/model_executor/model_loader/
- label: Multi-Modal Models (Extended Generation 1)
device: h200_35gb
key: multi-modal-models-extended-generation-1
optional: true
source_file_dependencies:
.buildkite/test_areas/pytorch.yaml
+38 -2
View File
@@ -116,8 +116,9 @@ steps:
- image-build-amd
- label: PyTorch Fullgraph Smoke Test
device: h200_35gb
key: pytorch-fullgraph-smoke-test
timeout_in_minutes: 60
timeout_in_minutes: 90
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
@@ -149,7 +150,42 @@ steps:
# as it is a heavy test that is covered in other steps.
# Use `find` to launch multiple instances of pytest so that
# they do not suffer from https://github.com/vllm-project/vllm/issues/28965
- "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
- "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
# Hopper-only DeepSeek-V2-Lite cases in this file require two 29.3-GiB model
# instances and cannot fit a 35GB MIG slice. L4 retains the original coverage:
# those SM90 cases skip while the architecture-compatible cases still run.
- label: PyTorch Fullgraph CUDAGraph (L4 Compatibility)
key: pytorch-fullgraph-cudagraph-l4-compatibility
timeout_in_minutes: 60
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
- vllm/_custom_ops.py
- vllm/compilation/
- vllm/config/
- vllm/distributed/
- vllm/engine/
- vllm/env_override.py
- vllm/envs.py
- vllm/forward_context.py
- vllm/inputs/
- vllm/ir/
- vllm/kernels/
- vllm/logger.py
- vllm/model_executor/
- vllm/multimodal/
- vllm/platforms/
- vllm/plugins/
- vllm/sampling_params.py
- vllm/sequence.py
- vllm/transformers_utils/
- vllm/triton_utils/
- vllm/utils/
- vllm/v1/
- tests/compile
commands:
- pytest -s -v compile/fullgraph/test_full_cudagraph.py
- label: PyTorch Fullgraph
key: pytorch-fullgraph
.buildkite/test_areas/quantization.yaml
+15 -4
View File
@@ -3,8 +3,11 @@ depends_on:
- image-build
steps:
- label: Quantization
device: h200_35gb
key: quantization
timeout_in_minutes: 60
timeout_in_minutes: 75
env:
VLLM_USE_V2_MODEL_RUNNER: "0"
source_file_dependencies:
- csrc/
- vllm/model_executor/layers/quantization
@@ -19,9 +22,13 @@ steps:
# TODO(jerryzh168): resolve the above comment
- uv pip install --system torchao==0.17.0 --index-url https://download.pytorch.org/whl/cu130
- uv pip install --system conch-triton-kernels
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
# The SM90-only checkpoint currently contains a removed weight_chan_scale
# parameter. It was not exercised by the previous L4 job.
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
parallelism: 8
- label: Quantized Fusions
device: h200_35gb
key: quantized-fusions
timeout_in_minutes: 20
source_file_dependencies:
@@ -52,10 +59,14 @@ steps:
- pytest -s -v tests/quantization/test_blackwell_moe.py
- label: Quantized Models Test
device: h200_35gb
key: quantized-models-test
timeout_in_minutes: 50
timeout_in_minutes: 65
env:
VLLM_USE_V2_MODEL_RUNNER: "0"
source_file_dependencies:
- vllm/model_executor/layers/quantization
- tests/models/quantization
commands:
- pytest -v -s models/quantization
- pytest -v -s models/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
parallelism: 3
.buildkite/test_areas/rust_frontend.yaml
+1
View File
@@ -81,6 +81,7 @@ steps:
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
- label: Rust Frontend Tool Use
device: h200_35gb
timeout_in_minutes: 25
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
cmake/cpu_extension.cmake
+3
View File
@@ -430,6 +430,7 @@ set(VLLM_EXT_SRC
"csrc/cpu/layernorm.cpp"
"csrc/cpu/mla_decode.cpp"
"csrc/cpu/pos_encoding.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/torch_bindings.cpp")
@@ -489,6 +490,7 @@ if (ENABLE_X86_ISA)
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/dnnl_kernels.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/cpu/torch_bindings.cpp"
# TODO: Remove these files
"csrc/cpu/activation.cpp"
@@ -502,6 +504,7 @@ if (ENABLE_X86_ISA)
"csrc/cpu/utils.cpp"
"csrc/cpu/spec_decode_utils.cpp"
"csrc/cpu/cpu_attn.cpp"
"csrc/cpu/mamba_cpu.cpp"
"csrc/cpu/dnnl_kernels.cpp"
"csrc/cpu/torch_bindings.cpp"
# TODO: Remove these files
cmake/external_projects/tml_fa4.cmake
+1 -1
View File
@@ -14,7 +14,7 @@ else()
FetchContent_Declare(
tml_fa4
GIT_REPOSITORY https://github.com/vllm-project/tml-fa4.git
GIT_TAG 13374f0c855acc1add1bf30444bd67aebbc24a8e
GIT_TAG b206834606ed5b5f21f8eed6b0683f528ea9cf7d
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND "")
csrc/cpu/cpu_types_vsx.hpp
+8 -7
View File
@@ -336,13 +336,14 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
reg.val[1] = fp16_to_fp32_bits(raw_lo);
}
float reduce_sum() const {
AliasReg ar;
ar.reg = reg;
float result = 0;
unroll_loop<int, VEC_ELEM_NUM>(
[&result, &ar](int i) { result += ar.values[i]; });
return result;
// VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
// Step 1: pairwise sum of the two 4-wide halves
__vector float s = vec_add(reg.val[0], reg.val[1]);
// Step 2: rotate by 8 bytes (2 floats) and add
s = vec_add(s, vec_sld(s, s, 8));
// Step 3: rotate by 4 bytes (1 float) and add => all lanes hold total
s = vec_add(s, vec_sld(s, s, 4));
return vec_extract(s, 0);
}
FP32Vec8 exp() const {
f32x4x2_t out;
csrc/cpu/mamba_cpu.cpp
+285
View File
@@ -0,0 +1,285 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// CPU at::Tensor wrappers for Mamba decode-step kernels defined in
// mamba_kernels.hpp.
#include "cpu/mamba_kernels.hpp"
#include <ATen/ATen.h>
#include <torch/library.h>
#include <c10/util/Optional.h>
#include "cpu_types.hpp"
// ---------------------------------------------------------------------------
// causal_conv1d_update
// ---------------------------------------------------------------------------
at::Tensor causal_conv1d_update_cpu_impl(
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
const c10::optional<at::Tensor>& bias,
const c10::optional<std::string>& activation,
const c10::optional<at::Tensor>& conv_state_indices,
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) {
bool do_silu = false;
if (activation.has_value()) {
const std::string& act = activation.value();
do_silu = (act == "silu" || act == "swish");
}
at::ScalarType dtype = x.scalar_type();
// Input x: contiguous in native dtype.
at::Tensor x_c = x.is_contiguous() ? x : x.contiguous();
// conv_state: NEVER copy the full paged tensor just for layout reasons.
// If the dtype matches we work directly on conv_state (contiguous or not)
// by extracting strides and passing them to the kernel.
// Only a dtype-conversion copy is made when types differ (rare for BF16).
bool state_type_ok = (conv_state.scalar_type() == dtype);
at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype);
// state_c and conv_state may be non-contiguous — that is intentional.
// Weight: coerce to same dtype if needed (should match in practice)
at::Tensor w_c =
(weight.scalar_type() != dtype)
? weight.to(dtype).contiguous()
: (weight.is_contiguous() ? weight : weight.contiguous());
// Bias stays float32 (small scalar, used only for fp32 accumulation)
at::Tensor bias_f32;
if (bias.has_value() && bias.value().defined())
bias_f32 = bias.value().to(at::kFloat).contiguous();
int64_t batch = x_c.size(0);
int64_t dim = x_c.size(1);
int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1;
int64_t width = w_c.size(1);
int64_t state_len = state_c.size(2);
// Extract strides — works for contiguous AND non-contiguous (transposed)
// state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in
// contiguous) stride(1): between conv channels (dim stride) stride(2):
// between state elements (=1 when contiguous, =dim when transposed)
int64_t stride_s_slot = state_c.stride(0);
int64_t stride_s_dim = state_c.stride(1);
int64_t stride_s_state = state_c.stride(2);
at::Tensor out = x_c.clone(); // native dtype, no float32 alloc
const int32_t* cache_idx_ptr = nullptr;
at::Tensor cache_idx_int;
if (conv_state_indices.has_value()) {
cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous();
cache_idx_ptr = cache_idx_int.data_ptr<int32_t>();
}
VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] {
mamba_cpu::causal_conv1d_update_kernel<scalar_t>(
x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot,
stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(),
bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr,
out.data_ptr<scalar_t>(), cache_idx_ptr,
static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len,
do_silu);
});
// Write back only when a type-conversion copy was made.
// Layout-only non-contiguity is handled via strides above — no copy needed.
if (!state_type_ok) conv_state.copy_(state_c);
return out;
}
// ---------------------------------------------------------------------------
// selective_state_update
// ---------------------------------------------------------------------------
void selective_state_update_cpu_impl(
at::Tensor& state, // (nstates, nheads, dim, dstate)
const at::Tensor& x, // (N, nheads, dim)
const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B,
const at::Tensor& C, const c10::optional<at::Tensor>& D,
const c10::optional<at::Tensor>& z,
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
const c10::optional<at::Tensor>& state_batch_indices,
const c10::optional<at::Tensor>& dst_state_batch_indices,
int64_t null_block_id, at::Tensor& out,
const c10::optional<at::Tensor>& num_accepted_tokens,
const c10::optional<at::Tensor>& cu_seqlens) {
at::ScalarType state_type = state.scalar_type();
at::ScalarType input_type = x.scalar_type();
// x, B, C must be contiguous and match input_type
auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor {
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor x_in = ensure_input(x);
at::Tensor B_in = ensure_input(B);
at::Tensor C_in = ensure_input(C);
at::Tensor z_in;
if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value());
// A, D, dt_bias are float32 model parameters that arrive here as expanded
// tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0).
// We need just the scalar value per head as a (nheads,) 1-D array so that
// A_ptr[h] in the kernel correctly reads head h's value.
//
// Strategy: peel trailing expanded (stride=0) dims via .select(), which is
// a zero-copy view. For A: (nheads, head_dim, dstate) strides (1,0,0)
// → .select(2,0) → (nheads, head_dim) strides (1,0)
// → .select(1,0) → (nheads,) stride (1,) ← contiguous, free.
// No allocation, no type conversion (A is already float32).
auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor {
at::Tensor r = t;
// Peel trailing dimensions that are broadcast (stride=0 or size=1)
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32
at::Tensor D_f32, dt_bias_f32;
if (D.has_value() && D.value().defined())
D_f32 = to_per_head_1d_f32(D.value());
if (dt_bias.has_value() && dt_bias.value().defined())
dt_bias_f32 = to_per_head_1d_f32(dt_bias.value());
// dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE
// the type conversion so we convert head_dim x fewer elements.
at::Tensor dt_f32;
{
// If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2,
// take a zero-copy view of index 0 along that dim first.
at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt; // (N, nheads)
at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2;
dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous();
}
int64_t nheads = state.size(1);
int64_t dim = state.size(2);
int64_t dstate = state.size(3);
int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined())
? cu_seqlens.value().size(0) - 1
: x_in.size(0);
int64_t ngroups = B_in.size(1);
// Strides
int64_t stride_state_n = state.stride(0);
int64_t stride_state_h = state.stride(1);
int64_t stride_state_d = state.stride(2);
int64_t stride_x_n = x_in.stride(0);
int64_t stride_x_h = x_in.stride(1);
int64_t stride_dt_n = dt_f32.stride(0); // dt is (N, nheads)
int64_t stride_BC_n = B_in.stride(0);
int64_t stride_BC_g = B_in.stride(1);
int64_t stride_out_n = out.stride(0);
int64_t stride_out_h = out.stride(1);
// Optional index pointers
auto get_int32_ptr =
[](const c10::optional<at::Tensor>& opt) -> const int32_t* {
return (opt.has_value() && opt.value().defined())
? opt.value().data_ptr<int32_t>()
: nullptr;
};
const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices);
const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices);
const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens);
const int32_t* csl_ptr = get_int32_ptr(cu_seqlens);
// Dispatch on (state_t, input_t, out_t): write directly into `out`
// without any intermediate float32 buffer.
VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] {
using state_t = scalar_t;
VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] {
using input_t = scalar_t;
VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] {
using out_t = scalar_t;
mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>(
state.data_ptr<state_t>(), stride_state_n, stride_state_h,
stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h,
dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(),
B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n,
stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
z_in.defined() ? z_in.data_ptr<input_t>() : nullptr,
dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr,
out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr,
dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N,
nheads, ngroups, dim, dstate, dt_softplus);
});
});
});
}
// ---------------------------------------------------------------------------
// mamba_chunk_scan_fwd_cpu
// ---------------------------------------------------------------------------
void mamba_chunk_scan_fwd_cpu_impl(
at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller
at::Tensor&
final_states, // [batch, nheads, headdim, dstate] float32 contiguous
const at::Tensor& x, // [seqlen, nheads, headdim]
const at::Tensor&
dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp)
const at::Tensor& A, // [nheads] float32
const at::Tensor& B, // [seqlen, ngroups, dstate]
const at::Tensor& C, // [seqlen, ngroups, dstate]
const c10::optional<at::Tensor>& D, // [nheads] float32 (optional)
const c10::optional<at::Tensor>& z, // [seqlen, nheads, headdim] (optional)
const at::Tensor& cu_seqlens // [batch+1] int32
) {
const at::ScalarType input_type = x.scalar_type();
auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor {
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor x_in = ensure_contig(x);
at::Tensor B_in = ensure_contig(B);
at::Tensor C_in = ensure_contig(C);
at::Tensor z_in;
if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value());
// A and D are float32 model parameters, potentially broadcast-expanded.
// Strip trailing broadcast dims to get a contiguous (nheads,) array.
auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor {
at::Tensor r = t;
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
return r.is_contiguous() ? r : r.contiguous();
};
at::Tensor A_f32 = to_per_head_f32(A);
at::Tensor D_f32;
if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value());
// dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in
// Python.
at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous();
if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat);
at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous();
const int64_t batch = final_states.size(0);
const int64_t nheads = final_states.size(1);
const int64_t headdim = final_states.size(2);
const int64_t dstate = final_states.size(3);
const int64_t ngroups = B_in.size(1);
TORCH_CHECK(final_states.is_contiguous(),
"mamba_chunk_scan_fwd_cpu: final_states must be contiguous");
TORCH_CHECK(out.is_contiguous(),
"mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via "
"raw data_ptr)");
VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] {
mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>(
final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(),
dt_c.data_ptr<float>(), A_f32.data_ptr<float>(),
B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(),
D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr,
out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads,
ngroups, headdim, dstate);
});
}
csrc/cpu/mamba_kernels.hpp
+382
View File
@@ -0,0 +1,382 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//
// Fused CPU vector kernels for Mamba decode-step hotspots:
// - causal_conv1d_update (depthwise 1-D conv state roll + compute)
// - selective_state_update (SSM recurrence, single-step)
#pragma once
#include "cpu_types.hpp"
#include <cmath>
#include <cstring>
#include <cstdint>
#include <algorithm>
namespace mamba_cpu {
// ---------------------------------------------------------------------------
// causal_conv1d_update — templated for native BF16/FP32
//
// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor.
// Explicit strides are passed so the kernel writes directly into the
// correct memory locations without making a contiguous copy of the full
// paged tensor (which was the source of the 34-41% direct_copy_kernel).
//
// stride_s_slot = state.stride(0) — between cache slots
// stride_s_dim = state.stride(1) — between conv_dim channels
// stride_s_state = state.stride(2) — between state elements
//
// When stride_s_state == 1 (contiguous), the memmove fast path is used.
// ---------------------------------------------------------------------------
template <typename scalar_t>
inline void causal_conv1d_update_kernel(
const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr,
int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state,
const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs,
int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen,
int64_t width, int64_t state_len, bool do_silu) {
#pragma omp parallel for
for (int64_t b = 0; b < batch; ++b) {
int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b;
if (cache_idx == pad_slot_id) continue;
for (int64_t t = 0; t < seqlen; ++t) {
const scalar_t* x_b = x_ptr + (b * dim * seqlen + t);
scalar_t* out_b = out_ptr + (b * dim * seqlen + t);
// Base of this slot in the (possibly non-contiguous) paged state
scalar_t* s_base = state_ptr + cache_idx * stride_s_slot;
for (int64_t d = 0; d < dim; ++d) {
float x_val = static_cast<float>(x_b[d * seqlen]);
scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state
const scalar_t* w = weight_ptr + d * width;
// Accumulate in float32 for precision
float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f;
for (int64_t k = 0; k < state_len; ++k) {
acc += static_cast<float>(w[k]) *
static_cast<float>(sd[k * stride_s_state]);
}
acc += static_cast<float>(w[state_len]) * x_val;
// Shift state left and append new input.
// Use memmove when contiguous (stride==1); element loop otherwise.
if (stride_s_state == 1) {
if (state_len > 1)
std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t));
if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val);
} else {
for (int64_t k = 0; k < state_len - 1; ++k)
sd[k * stride_s_state] = sd[(k + 1) * stride_s_state];
if (state_len > 0)
sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val);
}
if (do_silu) {
float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc))
: std::exp(acc) / (1.0f + std::exp(acc));
acc *= sigmoid;
}
out_b[d * seqlen] = static_cast<scalar_t>(acc);
}
}
}
}
// ---------------------------------------------------------------------------
// selective_state_update
//
// Template parameters:
// state_t - dtype of ssm_state cache (typically BFloat16)
// input_t - dtype of x, B, C (typically BFloat16)
// out_t - dtype of output tensor (typically BFloat16)
// Write directly — no float32 intermediate buffer needed.
//
// A, D, dt_bias are accepted as const float* (they are always float32
// model parameters in Mamba2). This eliminates the per-call float32→BF16
// conversion and the .contiguous() materialisation of the broadcast-expand.
//
// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the
// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed.
// ---------------------------------------------------------------------------
template <typename state_t, typename input_t, typename out_t = float>
inline void selective_state_update_kernel(
state_t* __restrict__ state_ptr, int64_t stride_state_n,
int64_t stride_state_h, int64_t stride_state_d,
const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h,
// dt: (N, nheads) — scalar per head, NOT expanded to head_dim
const float* __restrict__ dt_ptr, int64_t stride_dt_n,
// A: (nheads,) float32 — scalar per head
const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr,
const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g,
// D: (nheads,) float32 — scalar per head (nullptr if not used)
const float* __restrict__ D_ptr,
// z: same shape as x (optional)
const input_t* __restrict__ z_ptr,
// dt_bias: (nheads,) float32 — scalar per head (nullptr if not used)
const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr,
int64_t stride_out_n, int64_t stride_out_h,
const int32_t* __restrict__ state_batch_indices,
const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id,
const int32_t* __restrict__ num_accepted_tokens,
const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads,
int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) {
using state_vec_t = vec_op::vec_t<state_t>;
using input_vec_t = vec_op::vec_t<input_t>;
constexpr int VEC_ELEM_NUM = 8;
int64_t nheads_per_group = nheads / ngroups;
for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) {
int64_t bos, seq_len;
if (cu_seqlens != nullptr) {
bos = cu_seqlens[seq_idx];
seq_len = cu_seqlens[seq_idx + 1] - bos;
} else {
bos = seq_idx;
seq_len = 1;
}
int64_t state_read_idx = (state_batch_indices != nullptr)
? state_batch_indices[seq_idx]
: seq_idx;
if (state_read_idx == null_block_id) continue;
int64_t state_write_idx = (num_accepted_tokens == nullptr)
? ((dst_state_batch_indices != nullptr)
? dst_state_batch_indices[seq_idx]
: state_read_idx)
: -1;
state_t* s = state_ptr + state_read_idx * stride_state_n;
for (int64_t t = 0; t < seq_len; ++t) {
int64_t token_idx = bos + t;
const input_t* x_tok = x_ptr + token_idx * stride_x_n;
// dt: (N, nheads) — one float per head per token
const float* dt_tok = dt_ptr + token_idx * stride_dt_n;
const input_t* B_tok = B_ptr + token_idx * stride_BC_n;
const input_t* C_tok = C_ptr + token_idx * stride_BC_n;
out_t* out_tok = out_ptr + token_idx * stride_out_n;
#pragma omp parallel for
for (int64_t h = 0; h < nheads; ++h) {
int64_t g = h / nheads_per_group;
const input_t* x_h = x_tok + h * stride_x_h;
const input_t* B_g = B_tok + g * stride_BC_g;
const input_t* C_g = C_tok + g * stride_BC_g;
out_t* out_h = out_tok + h * stride_out_h;
state_t* s_h = s + h * stride_state_h;
// Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing
float dt_val = dt_tok[h];
if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h];
if (dt_softplus) {
dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val;
}
const float A_val = A_ptr[h]; // scalar: same for all dim, dstate
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
const input_t* z_h =
(z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h
: nullptr;
vec_op::FP32Vec8 dt_vec(dt_val);
// dA = exp(A * dt): A and dt are SCALARS per head, so compute once
// and broadcast. This saves 7 redundant std::exp() calls that
// FP32Vec8::exp() would otherwise make on the broadcast vector.
const float dA_scalar = std::exp(A_val * dt_val);
vec_op::FP32Vec8 dA(dA_scalar); // broadcast
for (int64_t d = 0; d < dim; ++d) {
float x_val = static_cast<float>(x_h[d]);
vec_op::FP32Vec8 out_vec(0.0f);
state_t* s_hd = s_h + d * stride_state_d;
const input_t* B_g_base = B_g;
const input_t* C_g_base = C_g;
vec_op::FP32Vec8 x_vec(x_val);
// dBx = B * x * dt — same dA for all dstate (A is scalar)
// s_new = s * dA + B * x * dt
int64_t n = 0;
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n)));
vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n)));
vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n)));
vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec;
vec_op::FP32Vec8 s_new = s_v * dA + dBx;
state_vec_t(s_new).save(s_hd + n);
out_vec = out_vec + s_new * C_v;
}
float out_val = out_vec.reduce_sum();
for (; n < dstate; ++n) {
// Reuse dA_scalar computed once per head — no exp() re-call
float dBx = static_cast<float>(B_g[n]) * x_val * dt_val;
float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx;
s_hd[n] = static_cast<state_t>(s_new);
out_val += s_new * static_cast<float>(C_g[n]);
}
if (D_ptr != nullptr) out_val += x_val * D_val;
if (z_h != nullptr) {
float z_val = static_cast<float>(z_h[d]);
float sigmoid = (z_val >= 0)
? 1.0f / (1.0f + std::exp(-z_val))
: std::exp(z_val) / (1.0f + std::exp(z_val));
out_val *= z_val * sigmoid;
}
out_h[d] = static_cast<out_t>(out_val);
}
}
if (num_accepted_tokens != nullptr &&
dst_state_batch_indices != nullptr) {
int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t];
if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) {
state_t* dst_s = state_ptr + token_dst_idx * stride_state_n;
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
}
}
}
if (num_accepted_tokens == nullptr && state_write_idx != null_block_id &&
state_write_idx != state_read_idx) {
state_t* dst_s = state_ptr + state_write_idx * stride_state_n;
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
}
}
}
// ---------------------------------------------------------------------------
// mamba_chunk_scan_fwd
//
// Prefill SSM recurrence for Mamba2 / SSD models.
//
// Key difference from selective_state_update_kernel (decode path):
// - #pragma omp parallel for collapse(2) is OUTSIDE the time loop.
// Each thread owns a (batch, head) slice and runs the entire token
// sequence without any per-token OpenMP synchronisation overhead.
// For seqlen=256, this eliminates 256 thread-barrier launches per batch.
//
// `dt` arrives already processed (float32, after bias + softplus + clamp)
// to keep this kernel simple. Preprocessing is done in the Python wrapper.
//
// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output
// tensor, pre-initialised by the caller (zero or from initial_states).
// Each (b, h) slice is private to exactly one thread via collapse(2), so
// there are no write conflicts.
//
// D is treated as a scalar per head ([nheads] float32).
// ---------------------------------------------------------------------------
template <typename input_t>
inline void mamba_chunk_scan_fwd_kernel(
float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32
const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim]
const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed)
const float* __restrict__ A_ptr, // [nheads] f32
const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate]
const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate]
const float* __restrict__ D_ptr, // [nheads] f32 (nullable)
const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable)
input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim]
const int32_t* __restrict__ cu_seqlens, // [batch+1] int32
int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim,
int64_t dstate) {
using input_vec_t = vec_op::vec_t<input_t>;
constexpr int VEC_ELEM_NUM = 8;
const int64_t nheads_per_group = nheads / ngroups;
// states layout: [batch, nheads, headdim, dstate] contiguous (caller
// guarantee)
const int64_t stride_s_b = nheads * headdim * dstate;
const int64_t stride_s_h = headdim * dstate;
// stride_s_d = dstate, stride_s_n = 1
#pragma omp parallel for collapse(2) schedule(static)
for (int64_t b = 0; b < batch; ++b) {
for (int64_t h = 0; h < nheads; ++h) {
const int64_t seq_start = cu_seqlens[b];
const int64_t seq_end = cu_seqlens[b + 1];
const int64_t g = h / nheads_per_group;
const float A_val = A_ptr[h];
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
// Working state slice: states[b, h, :, :] — float32, headdim * dstate.
// Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB).
float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h;
for (int64_t t = seq_start; t < seq_end; ++t) {
const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim;
const float* dt_h = dt_ptr + t * nheads + h;
const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate;
const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate;
const input_t* z_h = (z_ptr != nullptr)
? z_ptr + t * nheads * headdim + h * headdim
: nullptr;
input_t* out_h = out_ptr + t * nheads * headdim + h * headdim;
const float dt_val = *dt_h;
const float dA_val = std::exp(A_val * dt_val);
const vec_op::FP32Vec8 dA_vec(dA_val); // broadcast scalar
const vec_op::FP32Vec8 dt_vec(dt_val);
for (int64_t d = 0; d < headdim; ++d) {
const float x_val = static_cast<float>(x_h[d]);
float* s_bhd = s_bh + d * dstate; // [dstate] contiguous float32
// Vectorised SSM update + readout over dstate:
// s_new = s * dA + x * dt * B
// y += s_new * C
int64_t n = 0;
vec_op::FP32Vec8 y_vec(0.0f);
const vec_op::FP32Vec8 x_vec(x_val);
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n)));
const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n)));
const vec_op::FP32Vec8 s_v(s_bhd + n);
const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v;
s_new.save(s_bhd + n);
y_vec = y_vec + s_new * C_v;
}
float y_val = y_vec.reduce_sum();
// Scalar tail for remaining dstate elements
for (; n < dstate; ++n) {
const float B_n = static_cast<float>(B_g[n]);
const float C_n = static_cast<float>(C_g[n]);
const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n;
s_bhd[n] = s_new;
y_val += s_new * C_n;
}
// D skip connection (scalar per head)
if (D_ptr != nullptr) y_val += x_val * D_val;
// z gating: out = y * z * sigmoid(z) (SiLU)
if (z_h != nullptr) {
const float z_val = static_cast<float>(z_h[d]);
const float sigmoid =
(z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val))
: std::exp(z_val) / (1.0f + std::exp(z_val));
y_val *= z_val * sigmoid;
}
out_h[d] = static_cast<input_t>(y_val);
}
}
}
}
}
} // namespace mamba_cpu
csrc/cpu/torch_bindings.cpp
+50
View File
@@ -213,6 +213,32 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
torch::Tensor slot_mapping,
const int64_t block_size);
at::Tensor causal_conv1d_update_cpu_impl(
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
const c10::optional<at::Tensor>& bias,
const c10::optional<std::string>& activation,
const c10::optional<at::Tensor>& conv_state_indices,
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
void selective_state_update_cpu_impl(
at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
const c10::optional<at::Tensor>& state_batch_indices,
const c10::optional<at::Tensor>& dst_state_batch_indices,
int64_t null_block_id, at::Tensor& out,
const c10::optional<at::Tensor>& num_accepted_tokens,
const c10::optional<at::Tensor>& cu_seqlens);
void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
const at::Tensor& x, const at::Tensor& dt,
const at::Tensor& A, const at::Tensor& B,
const at::Tensor& C,
const c10::optional<at::Tensor>& D,
const c10::optional<at::Tensor>& z,
const at::Tensor& cu_seqlens);
void init_cpu_memory_env(std::vector<int64_t> node_ids);
namespace cpu_utils {
@@ -595,6 +621,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
"block_size) -> ()",
&compute_slot_mapping_kernel_impl);
// Mamba CPU kernels
ops.def(
"causal_conv1d_update_cpu_vec("
"Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
"Tensor? bias, str? activation, Tensor? conv_state_indices, "
"Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
&causal_conv1d_update_cpu_impl);
ops.def(
"selective_state_update_cpu("
"Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
"Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
"Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
"SymInt null_block_id, Tensor(a13!) out, "
"Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
&selective_state_update_cpu_impl);
ops.def(
"mamba_chunk_scan_fwd_cpu("
"Tensor(a0!) out, Tensor(a1!) final_states, "
"Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
"Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
&mamba_chunk_scan_fwd_cpu_impl);
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
// Speculative decoding kernels
csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu
+78
View File
@@ -71,6 +71,73 @@ __device__ __forceinline__ float toFloat(T value) {
}
}
#ifndef USE_ROCM
// Adapted from:
// https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh
template <typename OutIndType, typename HashIndType>
__launch_bounds__(128) __global__
void dsv4HashTopkSoftplusSqrt(const float* input, float* output,
OutIndType* indices, int num_rows,
int num_experts, float routed_scaling_factor,
const HashIndType* input_ids,
const HashIndType* tid2eid) {
const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int lane = threadIdx.x % 32;
if (warp >= num_rows) return;
const int64_t token_id = load_index_as_int64(input_ids, warp);
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaGridDependencySynchronize();
#endif
int expert = 0;
float weight = 0.f;
if (lane < 6) {
// only load and calculate for 6 experts
expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
const float x = input[warp * num_experts + expert];
weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
}
float weight_sum = weight;
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1) {
// sum in warp
weight_sum += VLLM_SHFL_XOR_SYNC(weight_sum, mask);
}
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaTriggerProgrammaticLaunchCompletion();
#endif
if (lane < 6) {
const int offset = warp * 6 + lane;
output[offset] =
weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);
indices[offset] = static_cast<OutIndType>(expert);
}
}
template <typename OutIndType, typename HashIndType>
void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
int num_rows, int num_experts,
double routed_scaling_factor,
const HashIndType* input_ids,
const HashIndType* tid2eid, cudaStream_t stream) {
if (num_rows == 0) return;
auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
cudaLaunchConfig_t config = {};
config.gridDim = (num_rows + 3) / 4;
config.blockDim = 128;
config.stream = stream;
cudaLaunchAttribute attr;
attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
attr.val.programmaticStreamSerializationAllowed = 1;
config.attrs = &attr;
config.numAttrs = 1;
const float scale = static_cast<float>(routed_scaling_factor);
cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
num_experts, scale, input_ids, tid2eid);
}
#endif
// ====================== TopK softplus_sqrt things
// ===============================
@@ -556,6 +623,17 @@ void topkGatingSoftplusSqrtKernelLauncher(
const float* correction_bias, const bool use_hash,
const HashIndType* input_ids, const HashIndType* tid2eid,
cudaStream_t stream) {
#ifndef USE_ROCM
if constexpr (std::is_same_v<InputType, float>) {
if (use_hash && topk == 6 && renormalize &&
(num_experts == 256 || num_experts == 384)) {
launchDsv4HashTopk<IndType, HashIndType>(
gating_output, topk_weights, topk_indices, num_tokens, num_experts,
routed_scaling_factor, input_ids, tid2eid, stream);
return;
}
}
#endif
static constexpr int WARPS_PER_TB = 4;
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
docs/design/fusions.md
+2 -2
View File
@@ -306,7 +306,7 @@ Supported quantization scheme/hardware combinations:
- Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py)
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu)
- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu)
### SiLU+Mul + Quantization (`fuse_act_quant`)
@@ -332,7 +332,7 @@ Supported quantization scheme/hardware combinations:
- Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py)
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
- CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/)
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
### RMSNorm + Padding (`fuse_act_padding`)
docs/features/kv_offloading_usage.md
+4 -3
View File
@@ -68,13 +68,14 @@ vllm serve <model> \
| --- | --- | --- | --- | --- |
| `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
## Secondary Tiers
@@ -179,7 +180,7 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
- For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
- `block_size` / `blocks_per_chunk`: larger offloaded chunks reduce per-block bookkeeping overhead but increase the granularity of lookups.
- FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
- Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
docs/models/hardware_supported_models/xpu.md
-2
View File
@@ -31,10 +31,8 @@
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| chuhac/TeleChat2-35B | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅ | | |
| 01-ai/Yi1.5-34B-Chat | YiForCausalLM | ✅ | | |
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| deepseek-ai/DeepSeek-Coder-33B-base | DeepSeekCoderForCausalLM | ✅ | | |
| meta-llama/Llama-2-13b-chat-hf | LlamaForCausalLM | ✅ | | |
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
| Qwen/Qwen1.5-14B-Chat | QwenForCausalLM | ✅ | | |
| Qwen/Qwen1.5-32B-Chat | QwenForCausalLM | ✅ | | |
| RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic | LlamaForCausalLM | | ✅ | |
examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py
+63 -13
View File
@@ -210,8 +210,31 @@ async def stream_decode_response(session, response, request_id):
await session.close()
def example_round_robin_dp_loader(request_number, dp_size):
return request_nums % dp_size
def flat_interleaved_dp_route(request_number, instances):
"""Flat round-robin over the full (instance, dp_rank) slot space.
ONE counter over (n_instances * dp_size) slots, so instance-selection and
DP-rank-selection are derived from the SAME index and can never alias. The
previous scheme computed instance = req % n and rank = req % dp from the
same counter with n | dp, which locked each instance to a stride-n subset
of its ranks (e.g. 2 prefill instances -> 4 of 8 ranks each -> half the
GPUs never receive a request, so the deployment falsely appears not to
scale).
Interleaved order — inst0_r0, inst1_r0, inst0_r1, inst1_r1, ... — so
consecutive requests alternate instances AND every rank gets walked.
Assumes homogeneous dp_size across a role's instances (true for the
DP<->DP and DP<->TP deployments this proxy targets). Returns
(instance_index, dp_rank); dp_rank is None when dp_size == 1 (e.g. a TP
decode), which avoids forwarding an out-of-range data-parallel rank.
"""
n = len(instances)
dp = instances[0]["dp_size"]
slot = (request_number - 1) % (n * dp)
inst_idx = slot % n
dp_rank = (slot // n) if dp > 1 else None
return inst_idx, dp_rank
@app.route("/health", methods=["GET"])
@@ -252,18 +275,21 @@ async def handle_request(api: str, request: Request):
503,
)
)
pid = request_nums % len(prefill_instances)
did = request_nums % len(decode_instances)
# Flat interleaved round-robin (see flat_interleaved_dp_route): ONE
# counter over the full (instance, dp_rank) slot space per role, so
# instance-selection and DP-rank-selection derive from the same index
# and can never alias. The old scheme keyed both on request_nums with
# n_instances | dp_size, stranding half the ranks (e.g. in 2P_DP8EP).
pid, selected_prefill_dp_rank = flat_interleaved_dp_route(
request_nums, prefill_instances
)
# Decode instance selection uses the same interleaved walk; in READ
# mode the decode reads KV from selected_prefill_dp_rank, so the
# decode's own dp_rank is not forwarded here.
did, _ = flat_interleaved_dp_route(request_nums, decode_instances)
prefill_instance_endpoint = prefill_instances[pid]
decode_instance_endpoint = decode_instances[did]
selected_prefill_dp_rank = None
if prefill_instance_endpoint["dp_size"] > 1:
selected_prefill_dp_rank = example_round_robin_dp_loader(
request_nums // len(prefill_instance_endpoint),
prefill_instance_endpoint["dp_size"],
)
# Embed both zmq_addresses in the request_id so the connector can parse
# the peer's host/ports from it, similar to P2P-NCCL
uid = str(uuid.uuid4()).replace("-", "")
@@ -427,9 +453,33 @@ if __name__ == "__main__":
args = parser.parse_args()
t = start_service_discovery("0.0.0.0", 36367)
app.debug = True
# High-concurrency hardening. Quart's app.run() uses a shallow listen
# backlog (100) and, with app.debug=True, adds per-request overhead that
# starves the single accept loop. Under a burst of ~512 simultaneous client
# connections the backlog overflows and the kernel RSTs the excess, so
# clients see "ClientOSError: [Errno 104] Connection reset by peer" before
# any response (~16% request loss at c=512). Serve via hypercorn with debug
# OFF and a deep backlog so the burst QUEUES (higher TTFT) instead of being
# reset -> 100% request success.
app.debug = False
app.config["BODY_TIMEOUT"] = 360000
app.config["RESPONSE_TIMEOUT"] = 360000
app.run(host="0.0.0.0", port=args.port)
import asyncio
import os
from hypercorn.asyncio import serve as _hypercorn_serve
from hypercorn.config import Config as _HypercornConfig
_hcfg = _HypercornConfig()
_hcfg.bind = [f"0.0.0.0:{args.port}"]
# Deep listen backlog so a wide connection burst queues, not RSTs. NOTE:
# effective backlog is capped by the host's net.core.somaxconn (proxy runs
# --network host); kernel 6.x defaults to 4096. Override via
# PROXY_LISTEN_BACKLOG.
_hcfg.backlog = int(os.environ.get("PROXY_LISTEN_BACKLOG", "4096"))
# Long-lived SSE streams (8k1k decode ~5 min): never reap on keepalive.
_hcfg.keep_alive_timeout = 360000.0
asyncio.run(_hypercorn_serve(app, _hcfg))
t.join()
examples/generate/multimodal/vision_language_multi_image_offline.py
+5
View File
@@ -17,6 +17,7 @@ from transformers import AutoProcessor, AutoTokenizer
from vllm import LLM, EngineArgs, SamplingParams
from vllm.lora.request import LoRARequest
from vllm.multimodal.utils import fetch_image
from vllm.platforms import current_platform
from vllm.utils.argparse_utils import FlexibleArgumentParser
QUESTION = "What is the content of each image?"
@@ -1443,6 +1444,8 @@ def run_generate(
engine_args.seed = seed
if tensor_parallel_size is not None:
engine_args.tensor_parallel_size = tensor_parallel_size
if current_platform.is_rocm():
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
llm = LLM.from_engine_args(engine_args)
sampling_params = SamplingParams(
@@ -1484,6 +1487,8 @@ def run_chat(
engine_args.seed = seed
if tensor_parallel_size is not None:
engine_args.tensor_parallel_size = tensor_parallel_size
if current_platform.is_rocm():
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
llm = LLM.from_engine_args(engine_args)
sampling_params = (
examples/generate/multimodal/vision_language_offline.py
+3
View File
@@ -21,6 +21,7 @@ from vllm.assets.image import ImageAsset
from vllm.assets.video import VideoAsset
from vllm.lora.request import LoRARequest
from vllm.multimodal.image import convert_image_mode
from vllm.platforms import current_platform
from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -2646,6 +2647,8 @@ def main(args):
if args.tensor_parallel_size is not None:
engine_args.tensor_parallel_size = args.tensor_parallel_size
engine_args = maybe_add_vit_cuda_graph_compilation_config(args, engine_args)
if current_platform.is_rocm():
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
llm = LLM.from_engine_args(engine_args)
# Don't want to check the flag multiple times, so just hijack `prompts`.
requirements/test/cpu.txt
+1 -1
View File
@@ -379,7 +379,7 @@ inflect==5.6.2
# via datamodel-code-generator
iniconfig==2.0.0
# via pytest
instanttensor==0.1.5
instanttensor==0.1.9
# via -r requirements/test/cuda.in
interegular==0.3.3
# via lm-format-enforcer
requirements/test/cuda.in
+1 -1
View File
@@ -58,7 +58,7 @@ arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix de
numba == 0.65.0 # Required for N-gram speculative decoding
runai-model-streamer[s3,gcs,azure]==0.15.7
fastsafetensors>=0.3.2
instanttensor>=0.1.5; platform_machine == "x86_64"
instanttensor>=0.1.9; platform_machine == "x86_64"
decord==0.6.0; platform_machine == "x86_64"
# terratorch is temporarily disabled while PyPI has the `lightning` package
# in `quarantined` status (every published terratorch version transitively
requirements/test/cuda.txt
+1 -1
View File
@@ -398,7 +398,7 @@ inflect==5.6.2
# via datamodel-code-generator
iniconfig==2.0.0
# via pytest
instanttensor==0.1.5
instanttensor==0.1.9
# via -r requirements/test/cuda.in
interegular==0.3.3
# via lm-format-enforcer
requirements/test/nightly-torch.txt
+1 -1
View File
@@ -44,5 +44,5 @@ numba == 0.65.0 # Required for N-gram speculative decoding
numpy
runai-model-streamer[s3,gcs,azure]==0.15.7
fastsafetensors>=0.3.2
instanttensor>=0.1.5
instanttensor>=0.1.9
pydantic>=2.12 # 2.11 leads to error on python 3.13
requirements/test/rocm.in
+1 -1
View File
@@ -54,7 +54,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
numba==0.65.0 # Required for N-gram speculative decoding
runai-model-streamer[s3,gcs,azure]==0.15.7
fastsafetensors>=0.3.2
instanttensor>=0.1.5
instanttensor>=0.1.9
decord==0.6.0
# Prithvi tests
requirements/test/rocm.txt
+1 -1
View File
@@ -391,7 +391,7 @@ inflect==7.5.0
# via datamodel-code-generator
iniconfig==2.3.0
# via pytest
instanttensor==0.1.6
instanttensor==0.1.9
# via -r requirements/test/rocm.in
interegular==0.3.3
# via lm-format-enforcer
rust/Cargo.lock
Generated
+1
View File
@@ -5560,6 +5560,7 @@ dependencies = [
"tracing",
"tracing-subscriber",
"uuid",
"vllm-bench",
"vllm-chat",
"vllm-engine-core-client",
"vllm-managed-engine",
rust/Cargo.toml
+1
View File
@@ -132,6 +132,7 @@ trait-set = "0.3.0"
url = "2.5.7"
uuid = { version = "1.22.0", features = ["v4"] }
validator = { version = "0.20.0", features = ["derive"] }
vllm-bench = { path = "src/bench" }
vllm-chat = { path = "src/chat" }
vllm-engine-core-client = { path = "src/engine-core-client" }
vllm-llm = { path = "src/llm" }
rust/src/bench/src/cli.rs
+4 -11
View File
@@ -3,8 +3,6 @@
use std::fmt;
use clap::Parser;
/// Backend type for the benchmark endpoint.
#[derive(clap::ValueEnum, Debug, Clone, Copy, PartialEq, Eq)]
pub enum BackendKind {
@@ -77,7 +75,7 @@ pub enum DatasetName {
ShareGpt,
#[value(name = "sonnet")]
Sonnet,
#[value(name = "speed-bench")]
#[value(name = "speed-bench", alias = "speed_bench")]
SpeedBench,
#[value(name = "hf")]
Hf,
@@ -144,13 +142,8 @@ impl fmt::Display for SpeedBenchConfig {
}
/// High-performance benchmark client for vLLM serving endpoints.
#[derive(Parser, Debug, Clone)]
#[command(
name = "vllm-bench",
about = "Benchmark online serving throughput",
version
)]
pub struct Cli {
#[derive(clap::Args, Debug, Clone)]
pub struct BenchServeArgs {
/// The type of backend or endpoint to use for the benchmark.
#[arg(long, default_value = "openai")]
pub backend: BackendKind,
@@ -659,7 +652,7 @@ pub struct Cli {
pub lora_assignment: LoraAssignment,
}
impl Cli {
impl BenchServeArgs {
/// Resolve the base URL from explicit --base-url or from --host/--port.
pub fn resolve_base_url(&self) -> String {
if let Some(ref base) = self.base_url {
rust/src/bench/src/config.rs
+212 -188
View File
@@ -4,7 +4,9 @@
use std::collections::HashMap;
use std::sync::Arc;
use crate::cli::{BackendKind, Cli, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig};
use crate::cli::{
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
};
use crate::datasets::random_mm::{MmBucketKey, MmLimitPerPrompt};
use crate::error::{BenchError, Result};
@@ -215,63 +217,63 @@ pub struct BenchConfig {
}
impl BenchConfig {
pub fn from_cli(cli: &Cli) -> Result<Self> {
if cli.burstiness <= 0.0 {
pub fn from_args(args: &BenchServeArgs) -> Result<Self> {
if args.burstiness <= 0.0 {
return Err(BenchError::Config("Burstiness must be positive".into()));
}
if cli.num_prompts == 0 {
if args.num_prompts == 0 {
return Err(BenchError::Config(
"--num-prompts must be at least 1".into(),
));
}
if cli.request_rate <= 0.0 && !cli.request_rate.is_infinite() {
if args.request_rate <= 0.0 && !args.request_rate.is_infinite() {
return Err(BenchError::Config(
"--request-rate must be positive (or inf)".into(),
));
}
if cli.max_model_len == Some(0) {
if args.max_model_len == Some(0) {
return Err(BenchError::Config(
"--max-model-len must be at least 1".into(),
));
}
let base_url = cli.resolve_base_url();
let api_url = cli.resolve_api_url();
let base_url = args.resolve_base_url();
let api_url = args.resolve_api_url();
let extra_headers = cli.parse_headers()?;
let mut extra_body = cli.parse_extra_body()?;
let extra_headers = args.parse_headers()?;
let mut extra_body = args.parse_extra_body()?;
// Merge sampling parameters into extra_body (matches Python behavior).
// Python collects non-None sampling params and merges them UNDER extra_body,
// meaning extra_body keys take precedence over sampling params.
{
let mut sampling_params = serde_json::Map::new();
if let Some(v) = cli.top_p {
if let Some(v) = args.top_p {
sampling_params.insert("top_p".into(), serde_json::json!(v));
}
if let Some(v) = cli.top_k {
if let Some(v) = args.top_k {
sampling_params.insert("top_k".into(), serde_json::json!(v));
}
if let Some(v) = cli.min_p {
if let Some(v) = args.min_p {
sampling_params.insert("min_p".into(), serde_json::json!(v));
}
if let Some(v) = cli.temperature {
if let Some(v) = args.temperature {
sampling_params.insert("temperature".into(), serde_json::json!(v));
}
if let Some(v) = cli.frequency_penalty {
if let Some(v) = args.frequency_penalty {
sampling_params.insert("frequency_penalty".into(), serde_json::json!(v));
}
if let Some(v) = cli.presence_penalty {
if let Some(v) = args.presence_penalty {
sampling_params.insert("presence_penalty".into(), serde_json::json!(v));
}
if let Some(v) = cli.repetition_penalty {
if let Some(v) = args.repetition_penalty {
sampling_params.insert("repetition_penalty".into(), serde_json::json!(v));
}
if !sampling_params.is_empty() {
if !cli.backend.is_openai_compatible() {
if !args.backend.is_openai_compatible() {
return Err(BenchError::Config(
"Sampling parameters are only supported by openai-compatible backends."
.into(),
@@ -299,7 +301,7 @@ impl BenchConfig {
}
// Parse metadata
let metadata = match &cli.metadata {
let metadata = match &args.metadata {
None => None,
Some(items) => {
let mut pairs = Vec::new();
@@ -314,24 +316,24 @@ impl BenchConfig {
};
// Parse goodput SLOs
let goodput = parse_goodput(&cli.goodput)?;
let goodput = parse_goodput(&args.goodput)?;
// Parse ramp-up config
let ramp_up = parse_ramp_up(cli)?;
let ramp_up = parse_ramp_up(args)?;
// Default percentile metrics based on backend type
let default_percentile_metrics = if cli.backend.is_pooling() {
let default_percentile_metrics = if args.backend.is_pooling() {
"e2el"
} else {
"ttft,tpot,itl,e2el"
};
let percentile_metrics_str =
cli.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
args.percentile_metrics.as_deref().unwrap_or(default_percentile_metrics);
let selected_percentile_metrics: Vec<String> =
percentile_metrics_str.split(',').map(|s| s.trim().to_string()).collect();
let metric_percentiles = parse_percentiles(&cli.metric_percentiles, false)?;
let sweep_summary_percentiles = cli
let metric_percentiles = parse_percentiles(&args.metric_percentiles, false)?;
let sweep_summary_percentiles = args
.sweep_summary_percentiles
.as_deref()
.map(|raw| parse_percentiles(raw, true))
@@ -344,38 +346,38 @@ impl BenchConfig {
selected_percentiles.push(90.0);
}
let tokenizer_id = if cli.skip_tokenizer_init {
let tokenizer_id = if args.skip_tokenizer_init {
None
} else {
Some(cli.tokenizer.clone().or_else(|| cli.model.clone()).unwrap_or_default())
args.tokenizer.clone().or_else(|| args.model.clone())
};
// Resolve input/output lengths
let random_input_len = cli.resolved_random_input_len();
let random_output_len = cli.resolved_random_output_len();
let per_turn_input_len = cli.resolved_per_turn_input_len();
let random_input_len = args.resolved_random_input_len();
let random_output_len = args.resolved_random_output_len();
let per_turn_input_len = args.resolved_per_turn_input_len();
// Normalized multi-turn turn counts (computed in validation block below, defaults
// to num_turns if multi-turn mode is not active)
let mut multi_turn_min_turns = cli.multi_turn_num_turns;
let mut multi_turn_max_turns = cli.multi_turn_num_turns;
let mut multi_turn_min_turns = args.multi_turn_num_turns;
let mut multi_turn_max_turns = args.multi_turn_num_turns;
// For random datasets with openai-compatible backends, default to ignore_eos.
// Exception: multi-turn mode, where ignore_eos causes unbounded context growth
// across turns. Multi-turn uses min_tokens instead for output length control.
// Pooling backends don't generate tokens, so ignore_eos is irrelevant.
let ignore_eos = if cli.backend.is_pooling() {
let ignore_eos = if args.backend.is_pooling() {
false
} else {
cli.ignore_eos
|| ((cli.dataset_name == DatasetName::Random
|| cli.dataset_name == DatasetName::RandomMm)
&& cli.backend.is_openai_compatible()
&& !cli.multi_turn)
args.ignore_eos
|| ((args.dataset_name == DatasetName::Random
|| args.dataset_name == DatasetName::RandomMm)
&& args.backend.is_openai_compatible()
&& !args.multi_turn)
};
// Pooling backends don't support multi-turn
if cli.backend.is_pooling() && cli.multi_turn {
if args.backend.is_pooling() && args.multi_turn {
return Err(BenchError::Config(
"Pooling/embedding backends do not support --multi-turn".into(),
));
@@ -383,7 +385,7 @@ impl BenchConfig {
// LoRA validation. Adapter names must be non-empty after trim; pooling
// backends are out of scope (vLLM LoRA routing is for generative paths).
let lora_modules = match cli.lora_modules.as_ref() {
let lora_modules = match args.lora_modules.as_ref() {
None => None,
Some(names) => {
if names.is_empty() {
@@ -391,7 +393,7 @@ impl BenchConfig {
"--lora-modules requires at least one adapter name".into(),
));
}
if cli.backend.is_pooling() {
if args.backend.is_pooling() {
return Err(BenchError::Config(
"--lora-modules is not supported for pooling/embedding backends".into(),
));
@@ -411,18 +413,18 @@ impl BenchConfig {
};
// Random-MM validation and config parsing
let (random_mm_limit, random_mm_buckets) = if cli.dataset_name == DatasetName::RandomMm {
if cli.backend != BackendKind::OpenaiChat {
let (random_mm_limit, random_mm_buckets) = if args.dataset_name == DatasetName::RandomMm {
if args.backend != BackendKind::OpenaiChat {
return Err(BenchError::Config(
"Multi-modal content (images) is only supported on 'openai-chat' backend."
.into(),
));
}
let limit = crate::datasets::random_mm::parse_limit_mm_per_prompt(
&cli.random_mm_limit_mm_per_prompt,
&args.random_mm_limit_mm_per_prompt,
)?;
let buckets =
crate::datasets::random_mm::parse_bucket_config(&cli.random_mm_bucket_config)?;
crate::datasets::random_mm::parse_bucket_config(&args.random_mm_bucket_config)?;
(limit, buckets)
} else {
(MmLimitPerPrompt::default(), Vec::new())
@@ -432,18 +434,18 @@ impl BenchConfig {
// sonnet (uses built-in Shakespeare's sonnets).
// Range ratio (Python semantics: [len*(1-r), len*(1+r)], each r in [0,1))
let random_range_ratio = RangeRatio::parse(&cli.random_range_ratio)?;
let random_range_ratio = RangeRatio::parse(&args.random_range_ratio)?;
// Batched inputs only make sense for pooling backends (the generation
// backends send one prompt per request).
if cli.random_batch_size == 0 {
if args.random_batch_size == 0 {
return Err(BenchError::Config(
"--random-batch-size must be at least 1".into(),
));
}
if cli.random_batch_size > 1
&& !cli.backend.is_pooling()
&& cli.dataset_name != DatasetName::RandomRerank
if args.random_batch_size > 1
&& !args.backend.is_pooling()
&& args.dataset_name != DatasetName::RandomRerank
{
return Err(BenchError::Config(
"--random-batch-size > 1 is only supported with embeddings/pooling backends".into(),
@@ -451,16 +453,16 @@ impl BenchConfig {
}
// random-rerank validation (mirrors Python RandomDatasetForReranking)
let is_reranker = !cli.no_reranker;
if cli.dataset_name == DatasetName::RandomRerank {
if !cli.backend.is_pooling() {
let is_reranker = !args.no_reranker;
if args.dataset_name == DatasetName::RandomRerank {
if !args.backend.is_pooling() {
return Err(BenchError::Config(
"--dataset-name random-rerank requires an embeddings/pooling backend \
(e.g. --backend vllm-rerank)"
.into(),
));
}
if !is_reranker && (cli.num_prompts < 2 || cli.random_batch_size < 2) {
if !is_reranker && (args.num_prompts < 2 || args.random_batch_size < 2) {
return Err(BenchError::Config(
"--no-reranker requires --num-prompts > 1 and --random-batch-size > 1 \
(the query is folded into the first batch slot)"
@@ -470,8 +472,8 @@ impl BenchConfig {
}
// Custom dataset validation
if cli.dataset_name == DatasetName::Custom {
match cli.dataset_path.as_deref() {
if args.dataset_name == DatasetName::Custom {
match args.dataset_path.as_deref() {
None => {
return Err(BenchError::Config(
"--dataset-path is required for --dataset-name custom \
@@ -486,7 +488,7 @@ impl BenchConfig {
}
_ => {}
}
if !cli.skip_chat_template {
if !args.skip_chat_template {
eprintln!(
"NOTE: client-side chat template rendering is not supported; custom \
dataset prompts are sent raw (equivalent to --skip-chat-template)."
@@ -495,29 +497,29 @@ impl BenchConfig {
}
// Prefix repetition validation
if cli.dataset_name == DatasetName::PrefixRepetition {
if cli.prefix_repetition_num_prefixes == 0 {
if args.dataset_name == DatasetName::PrefixRepetition {
if args.prefix_repetition_num_prefixes == 0 {
return Err(BenchError::Config(
"--prefix-repetition-num-prefixes must be at least 1".into(),
));
}
if cli.num_prompts < cli.prefix_repetition_num_prefixes {
if args.num_prompts < args.prefix_repetition_num_prefixes {
return Err(BenchError::Config(format!(
"--num-prompts ({}) must be >= --prefix-repetition-num-prefixes ({})",
cli.num_prompts, cli.prefix_repetition_num_prefixes
args.num_prompts, args.prefix_repetition_num_prefixes
)));
}
}
// HF dataset validation
if cli.dataset_name == DatasetName::Hf && cli.dataset_path.is_none() {
if args.dataset_name == DatasetName::Hf && args.dataset_path.is_none() {
return Err(BenchError::Config(
"--dataset-path is required for --dataset-name hf \
(set to a HuggingFace dataset ID, e.g. 'allenai/WildChat-4.8M')"
.into(),
));
}
if let Some(len) = cli.hf_output_len
if let Some(len) = args.hf_output_len
&& len == 0
{
return Err(BenchError::Config(
@@ -526,13 +528,13 @@ impl BenchConfig {
}
// Multi-turn validation
if cli.multi_turn {
if cli.backend != BackendKind::OpenaiChat {
if args.multi_turn {
if args.backend != BackendKind::OpenaiChat {
return Err(BenchError::Config(
"--multi-turn requires --backend openai-chat".into(),
));
}
if cli.multi_turn_num_turns == 0 {
if args.multi_turn_num_turns == 0 {
return Err(BenchError::Config(
"--multi-turn-num-turns must be at least 1".into(),
));
@@ -541,18 +543,18 @@ impl BenchConfig {
// Normalize and validate min/max turns. ShareGPT only consumes max_turns
// (the loader walks all available turns up to the cap), so the
// min/num/max coupling used for synthetic generation does not apply.
if cli.dataset_name == DatasetName::ShareGpt {
if cli.multi_turn_max_turns == 1 {
if args.dataset_name == DatasetName::ShareGpt {
if args.multi_turn_max_turns == 1 {
return Err(BenchError::Config(
"--multi-turn-max-turns must be at least 2 for ShareGPT multi-turn".into(),
));
}
} else {
(multi_turn_min_turns, multi_turn_max_turns) =
match (cli.multi_turn_min_turns, cli.multi_turn_max_turns) {
(0, 0) => (cli.multi_turn_num_turns, cli.multi_turn_num_turns),
(m, 0) => (m, cli.multi_turn_num_turns),
(0, x) => (cli.multi_turn_num_turns, x),
match (args.multi_turn_min_turns, args.multi_turn_max_turns) {
(0, 0) => (args.multi_turn_num_turns, args.multi_turn_num_turns),
(m, 0) => (m, args.multi_turn_num_turns),
(0, x) => (args.multi_turn_num_turns, x),
(m, x) => (m, x),
};
if multi_turn_min_turns < 1 {
@@ -575,8 +577,8 @@ impl BenchConfig {
}
// Validate prefix sharing ratios
let pg = cli.multi_turn_prefix_global_ratio;
let pc = cli.multi_turn_prefix_conversation_ratio;
let pg = args.multi_turn_prefix_global_ratio;
let pc = args.multi_turn_prefix_conversation_ratio;
if !(0.0..=1.0).contains(&pg) {
return Err(BenchError::Config(
"--multi-turn-prefix-global-ratio must be in [0.0, 1.0]".into(),
@@ -592,20 +594,20 @@ impl BenchConfig {
"--multi-turn-prefix-global-ratio + --multi-turn-prefix-conversation-ratio must be < 1.0 (unique suffix required)".into(),
));
}
if (pg > 0.0 || pc > 0.0) && cli.dataset_name != DatasetName::Random {
if (pg > 0.0 || pc > 0.0) && args.dataset_name != DatasetName::Random {
return Err(BenchError::Config(
"Prefix sharing (--multi-turn-prefix-global-ratio / --multi-turn-prefix-conversation-ratio) only works with --dataset-name random".into(),
));
}
}
if !(cli.steady_state_threshold > 0.0 && cli.steady_state_threshold <= 1.0) {
if !(args.steady_state_threshold > 0.0 && args.steady_state_threshold <= 1.0) {
return Err(BenchError::Config(format!(
"--steady-state-threshold must be in (0.0, 1.0], got {}",
cli.steady_state_threshold
args.steady_state_threshold
)));
}
if let Some(mw) = cli.steady_state_min_window
if let Some(mw) = args.steady_state_min_window
&& mw < 0.0
{
return Err(BenchError::Config(format!(
@@ -613,122 +615,122 @@ impl BenchConfig {
)));
}
if cli.profile_batch_threshold.is_some() && !cli.profile {
if args.profile_batch_threshold.is_some() && !args.profile {
return Err(BenchError::Config(
"--profile-batch-threshold requires --profile".into(),
));
}
if cli.profile_duration <= 0.0 {
if args.profile_duration <= 0.0 {
return Err(BenchError::Config(
"--profile-duration must be positive".into(),
));
}
if cli.profile_batch_threshold.is_none() && cli.profile_duration != 5.0 {
if args.profile_batch_threshold.is_none() && args.profile_duration != 5.0 {
return Err(BenchError::Config(
"--profile-duration requires --profile-batch-threshold".into(),
));
}
Ok(BenchConfig {
backend: cli.backend,
backend: args.backend,
base_url,
api_url,
model: cli.model.clone(),
model_name: cli.served_model_name.clone(),
model: args.model.clone(),
model_name: args.served_model_name.clone(),
tokenizer_id,
tokenizer_mode: cli.tokenizer_mode.clone(),
trust_remote_code: cli.trust_remote_code,
skip_tokenizer_init: cli.skip_tokenizer_init,
dataset_name: cli.dataset_name,
dataset_path: cli.dataset_path.clone(),
max_model_len: cli.max_model_len,
tokenizer_mode: args.tokenizer_mode.clone(),
trust_remote_code: args.trust_remote_code,
skip_tokenizer_init: args.skip_tokenizer_init,
dataset_name: args.dataset_name,
dataset_path: args.dataset_path.clone(),
max_model_len: args.max_model_len,
random_input_len,
random_output_len,
random_prefix_len: cli.random_prefix_len,
random_prefix_len: args.random_prefix_len,
random_range_ratio,
random_batch_size: cli.random_batch_size,
random_batch_size: args.random_batch_size,
is_reranker,
custom_output_len: cli.output_len.map(|v| v as i64).unwrap_or(cli.custom_output_len),
prefix_repetition_prefix_len: cli.prefix_repetition_prefix_len,
prefix_repetition_suffix_len: cli.prefix_repetition_suffix_len,
prefix_repetition_num_prefixes: cli.prefix_repetition_num_prefixes,
prefix_repetition_output_len: cli
custom_output_len: args.output_len.map(|v| v as i64).unwrap_or(args.custom_output_len),
prefix_repetition_prefix_len: args.prefix_repetition_prefix_len,
prefix_repetition_suffix_len: args.prefix_repetition_suffix_len,
prefix_repetition_num_prefixes: args.prefix_repetition_num_prefixes,
prefix_repetition_output_len: args
.output_len
.unwrap_or(cli.prefix_repetition_output_len),
random_cache_hit_fraction: cli.random_cache_hit_fraction,
random_cache_ratio: cli.random_cache_ratio,
sharegpt_output_len: cli.sharegpt_output_len,
sonnet_input_len: cli.sonnet_input_len,
sonnet_output_len: cli.sonnet_output_len,
sonnet_prefix_len: cli.sonnet_prefix_len,
no_oversample: cli.no_oversample,
disable_shuffle: cli.disable_shuffle,
num_prompts: cli.num_prompts,
request_rate: cli.request_rate,
burstiness: cli.burstiness,
max_concurrency: cli.max_concurrency,
steady_state_threshold: cli.steady_state_threshold,
steady_state_min_window: cli.steady_state_min_window,
no_steady_state: cli.no_steady_state,
disable_tqdm: cli.disable_tqdm,
num_warmups: cli.num_warmups,
profile: cli.profile,
profile_batch_threshold: cli.profile_batch_threshold,
profile_duration: cli.profile_duration,
save_result: cli.save_result,
save_detailed: cli.save_detailed,
append_result: cli.append_result,
result_dir: cli.result_dir.clone(),
result_filename: cli.result_filename.clone(),
seed: cli.seed,
.unwrap_or(args.prefix_repetition_output_len),
random_cache_hit_fraction: args.random_cache_hit_fraction,
random_cache_ratio: args.random_cache_ratio,
sharegpt_output_len: args.sharegpt_output_len,
sonnet_input_len: args.sonnet_input_len,
sonnet_output_len: args.sonnet_output_len,
sonnet_prefix_len: args.sonnet_prefix_len,
no_oversample: args.no_oversample,
disable_shuffle: args.disable_shuffle,
num_prompts: args.num_prompts,
request_rate: args.request_rate,
burstiness: args.burstiness,
max_concurrency: args.max_concurrency,
steady_state_threshold: args.steady_state_threshold,
steady_state_min_window: args.steady_state_min_window,
no_steady_state: args.no_steady_state,
disable_tqdm: args.disable_tqdm,
num_warmups: args.num_warmups,
profile: args.profile,
profile_batch_threshold: args.profile_batch_threshold,
profile_duration: args.profile_duration,
save_result: args.save_result,
save_detailed: args.save_detailed,
append_result: args.append_result,
result_dir: args.result_dir.clone(),
result_filename: args.result_filename.clone(),
seed: args.seed,
ignore_eos,
insecure: cli.insecure,
insecure: args.insecure,
selected_percentile_metrics,
selected_percentiles,
sweep_summary_percentiles,
label: cli.label.clone(),
logprobs: cli.logprobs,
request_id_prefix: cli.get_request_id_prefix(),
ready_check_timeout_sec: cli.ready_check_timeout_sec,
label: args.label.clone(),
logprobs: args.logprobs,
request_id_prefix: args.get_request_id_prefix(),
ready_check_timeout_sec: args.ready_check_timeout_sec,
extra_headers,
extra_body,
metadata,
dry_run: cli.dry_run,
dry_run: args.dry_run,
goodput,
ramp_up,
multi_turn: cli.multi_turn,
multi_turn_num_turns: cli.multi_turn_num_turns,
multi_turn: args.multi_turn,
multi_turn_num_turns: args.multi_turn_num_turns,
multi_turn_min_turns,
multi_turn_max_turns,
sharegpt_multi_turn_max_turns: if cli.multi_turn
&& cli.dataset_name == DatasetName::ShareGpt
&& cli.multi_turn_max_turns != 0
sharegpt_multi_turn_max_turns: if args.multi_turn
&& args.dataset_name == DatasetName::ShareGpt
&& args.multi_turn_max_turns != 0
{
Some(cli.multi_turn_max_turns)
Some(args.multi_turn_max_turns)
} else {
None
},
per_turn_input_len,
multi_turn_concurrency: cli.multi_turn_concurrency,
multi_turn_delay_ms: cli.multi_turn_delay_ms,
multi_turn_prefix_global_ratio: cli.multi_turn_prefix_global_ratio,
multi_turn_prefix_conversation_ratio: cli.multi_turn_prefix_conversation_ratio,
speed_bench_config: cli.speed_bench_config,
speed_bench_category: cli.speed_bench_category.clone(),
speed_bench_max_input_len: cli.speed_bench_max_input_len,
hf_split: cli.hf_split.clone(),
hf_subset: cli.hf_subset.clone(),
hf_output_len: cli.hf_output_len,
hf_text_column: cli.hf_text_column.clone(),
reset_prefix_cache: cli.reset_prefix_cache,
prompt_token_ids: cli.prompt_token_ids,
random_mm_base_items_per_request: cli.random_mm_base_items_per_request,
random_mm_num_mm_items_range_ratio: cli.random_mm_num_mm_items_range_ratio,
multi_turn_concurrency: args.multi_turn_concurrency,
multi_turn_delay_ms: args.multi_turn_delay_ms,
multi_turn_prefix_global_ratio: args.multi_turn_prefix_global_ratio,
multi_turn_prefix_conversation_ratio: args.multi_turn_prefix_conversation_ratio,
speed_bench_config: args.speed_bench_config,
speed_bench_category: args.speed_bench_category.clone(),
speed_bench_max_input_len: args.speed_bench_max_input_len,
hf_split: args.hf_split.clone(),
hf_subset: args.hf_subset.clone(),
hf_output_len: args.hf_output_len,
hf_text_column: args.hf_text_column.clone(),
reset_prefix_cache: args.reset_prefix_cache,
prompt_token_ids: args.prompt_token_ids,
random_mm_base_items_per_request: args.random_mm_base_items_per_request,
random_mm_num_mm_items_range_ratio: args.random_mm_num_mm_items_range_ratio,
random_mm_limit,
random_mm_buckets,
enable_multimodal_chat: cli.enable_multimodal_chat,
enable_multimodal_chat: args.enable_multimodal_chat,
lora_modules,
lora_assignment: cli.lora_assignment,
lora_assignment: args.lora_assignment,
})
}
}
@@ -811,17 +813,17 @@ fn parse_goodput(goodput_args: &Option<Vec<String>>) -> Result<GoodputConfig> {
Ok(config)
}
fn parse_ramp_up(cli: &Cli) -> Result<Option<RampUpConfig>> {
let strategy = match cli.ramp_up_strategy {
fn parse_ramp_up(args: &BenchServeArgs) -> Result<Option<RampUpConfig>> {
let strategy = match args.ramp_up_strategy {
None => return Ok(None),
Some(s) => s,
};
let start_rps = cli.ramp_up_start_rps.ok_or_else(|| {
let start_rps = args.ramp_up_start_rps.ok_or_else(|| {
BenchError::Config("--ramp-up-start-rps is required when --ramp-up-strategy is set".into())
})?;
let end_rps = cli.ramp_up_end_rps.ok_or_else(|| {
let end_rps = args.ramp_up_end_rps.ok_or_else(|| {
BenchError::Config("--ramp-up-end-rps is required when --ramp-up-strategy is set".into())
})?;
@@ -843,7 +845,21 @@ mod tests {
use clap::Parser;
use super::*;
use crate::cli::Cli;
use crate::cli::BenchServeArgs;
#[derive(Parser)]
struct TestCli {
#[command(flatten)]
args: BenchServeArgs,
}
fn parse_args<I, T>(args: I) -> BenchServeArgs
where
I: IntoIterator<Item = T>,
T: Into<std::ffi::OsString> + Clone,
{
TestCli::parse_from(args).args
}
fn base_multi_turn_args() -> Vec<&'static str> {
vec![
@@ -859,8 +875,8 @@ mod tests {
#[test]
fn test_prefix_sharing_defaults_to_zero() {
let args = base_multi_turn_args();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.multi_turn_prefix_global_ratio, 0.0);
assert_eq!(config.multi_turn_prefix_conversation_ratio, 0.0);
}
@@ -874,8 +890,8 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.8",
]);
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert!((config.multi_turn_prefix_global_ratio - 0.1).abs() < 1e-10);
assert!((config.multi_turn_prefix_conversation_ratio - 0.8).abs() < 1e-10);
}
@@ -889,8 +905,8 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.6",
]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -902,16 +918,16 @@ mod tests {
"--multi-turn-prefix-conversation-ratio",
"0.5",
]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
fn test_prefix_sharing_out_of_range_fails() {
let mut args = base_multi_turn_args();
args.extend(["--multi-turn-prefix-global-ratio", "1.5"]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -928,8 +944,8 @@ mod tests {
"--multi-turn-prefix-global-ratio",
"0.1",
];
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -944,8 +960,8 @@ mod tests {
"--dataset-name",
"sharegpt",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.multi_turn_max_turns, 3);
assert_eq!(config.sharegpt_multi_turn_max_turns, None);
@@ -968,8 +984,8 @@ mod tests {
"--multi-turn-max-turns",
"2",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(2));
}
@@ -987,8 +1003,8 @@ mod tests {
"--multi-turn-max-turns",
"1",
];
let cli = Cli::parse_from(args);
let err = BenchConfig::from_cli(&cli).unwrap_err().to_string();
let args = parse_args(args);
let err = BenchConfig::from_args(&args).unwrap_err().to_string();
assert!(
err.contains("at least 2 for ShareGPT"),
"expected ShareGPT-specific error, got: {err}"
@@ -1009,8 +1025,8 @@ mod tests {
"--multi-turn-max-turns",
"20",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sharegpt_multi_turn_max_turns, Some(20));
}
@@ -1018,8 +1034,8 @@ mod tests {
#[test]
fn test_sweep_summary_percentiles_default_empty() {
let args = base_multi_turn_args();
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert!(config.sweep_summary_percentiles.is_empty());
assert_eq!(config.selected_percentiles, vec![99.0, 90.0]);
@@ -1034,8 +1050,8 @@ mod tests {
"--sweep-summary-percentiles",
"90,95,90",
]);
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.sweep_summary_percentiles, vec![90.0, 95.0]);
assert_eq!(config.selected_percentiles, vec![99.0, 95.0, 90.0]);
@@ -1045,8 +1061,8 @@ mod tests {
fn test_invalid_sweep_summary_percentile_fails() {
let mut args = base_multi_turn_args();
args.extend(["--sweep-summary-percentiles", "101"]);
let cli = Cli::parse_from(args);
assert!(BenchConfig::from_cli(&cli).is_err());
let args = parse_args(args);
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
@@ -1058,12 +1074,20 @@ mod tests {
"--max-model-len",
"4096",
];
let cli = Cli::parse_from(args);
let config = BenchConfig::from_cli(&cli).unwrap();
let args = parse_args(args);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.max_model_len, Some(4096));
}
#[test]
fn test_tokenizer_id_deferred_when_model_is_unspecified() {
let args = parse_args(["vllm-bench"]);
let config = BenchConfig::from_args(&args).unwrap();
assert_eq!(config.tokenizer_id, None);
}
#[test]
fn test_zero_max_model_len_fails() {
let args = vec![
@@ -1073,9 +1097,9 @@ mod tests {
"--max-model-len",
"0",
];
let cli = Cli::parse_from(args);
let args = parse_args(args);
assert!(BenchConfig::from_cli(&cli).is_err());
assert!(BenchConfig::from_args(&args).is_err());
}
#[test]
fn test_range_ratio_parse_float() {
rust/src/bench/src/hub.rs
+5 -2
View File
@@ -40,8 +40,11 @@ impl HubRepo {
.build()
.map_err(|e| format!("Failed to build download runtime: {e}"))?;
rt.block_on(async move {
let api = hf_hub::api::tokio::Api::new()
.map_err(|e| format!("Failed to init HF API: {e}"))?;
let mut builder = hf_hub::api::tokio::ApiBuilder::from_env();
if let Ok(token) = std::env::var("HF_TOKEN") {
builder = builder.with_token(Some(token));
}
let api = builder.build().map_err(|e| format!("Failed to init HF API: {e}"))?;
api.repo(repo).get(&filename).await.map_err(|e| format!("{e}"))
})
})
rust/src/bench/src/lib.rs
+86
View File
@@ -0,0 +1,86 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
mod backends;
mod benchmark;
mod cli;
mod compare;
mod config;
mod datasets;
mod error;
mod hub;
mod metrics;
mod multi_run;
mod multi_turn;
mod output;
mod rate_control;
mod ready_checker;
mod sweep;
mod tiktoken;
mod tokenizer;
use anyhow::Context;
pub use cli::{
BackendKind, BenchServeArgs, DatasetName, LoraAssignment, RampUpStrategy, SpeedBenchConfig,
};
use config::BenchConfig;
/// Prepare process-wide resources for a benchmark run.
pub fn prepare_process() {
// Raise the open-file soft limit to the hard limit. High-concurrency
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
eprintln!("Open-file limit: {new}");
}
}
/// Run the online serving benchmark.
pub async fn run(args: BenchServeArgs) -> anyhow::Result<()> {
// --- Compare mode: no server needed, just diff two JSON files ---
if let Some(ref files) = args.compare {
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
}
let config = BenchConfig::from_args(&args).context("Configuration error")?;
async {
if config.multi_turn {
if let Some(ref sweep_mc) = args.sweep_max_concurrency {
// --- Sweep over concurrency in multi-turn mode ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_multi_turn_concurrency_sweep(
&config,
&values,
args.sweep_num_prompts_factor,
)
.await?;
} else {
// --- Single multi-turn conversation benchmark ---
multi_turn::run_multi_turn_benchmark(&config).await?;
}
} else if let Some(ref sweep_mc) = args.sweep_max_concurrency {
// --- Sweep over max-concurrency ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_concurrency_sweep(&config, &values, args.sweep_num_prompts_factor).await?;
} else if let Some(ref sweep_rate) = args.sweep_request_rate {
// --- Sweep over request-rate ---
let values =
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
sweep::run_rate_sweep(&config, &values).await?;
} else if args.num_runs > 1 {
// --- Multi-run with statistical aggregation ---
multi_run::run_multi(&config, args.num_runs).await?;
} else {
// --- Normal single benchmark ---
benchmark::run_benchmark(&config).await?;
}
anyhow::Ok(())
}
.await
.context("Benchmark failed")
}
rust/src/bench/src/main.rs
+14 -74
View File
@@ -1,92 +1,32 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
mod backends;
mod benchmark;
mod cli;
mod compare;
mod config;
mod datasets;
mod error;
mod hub;
mod metrics;
mod multi_run;
mod multi_turn;
mod output;
mod rate_control;
mod ready_checker;
mod sweep;
mod tiktoken;
mod tokenizer;
#[cfg(not(target_env = "msvc"))]
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
use anyhow::Context;
use clap::Parser;
use cli::Cli;
use config::BenchConfig;
#[derive(Parser)]
#[command(
name = "vllm-bench",
about = "Benchmark online serving throughput",
version
)]
struct Cli {
#[command(flatten)]
args: vllm_bench::BenchServeArgs,
}
fn main() -> anyhow::Result<()> {
// Raise the open-file soft limit to the hard limit. High-concurrency
// benchmarks (1024+ requests) easily exceed the default 1024 fd soft limit.
if let Ok(new) = rlimit::increase_nofile_limit(u64::MAX)
&& new > 1024
{
eprintln!("Open-file limit: {new}");
}
let cli = Cli::parse();
// --- Compare mode: no server needed, just diff two JSON files ---
if let Some(ref files) = cli.compare {
return compare::compare_results(&files[0], &files[1]).context("Comparison failed");
}
let config = BenchConfig::from_cli(&cli).context("Configuration error")?;
vllm_bench::prepare_process();
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.expect("Failed to build tokio runtime");
.context("Failed to build tokio runtime")?;
runtime
.block_on(async {
if config.multi_turn {
if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
// --- Sweep over concurrency in multi-turn mode ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_multi_turn_concurrency_sweep(
&config,
&values,
cli.sweep_num_prompts_factor,
)
.await?;
} else {
// --- Single multi-turn conversation benchmark ---
multi_turn::run_multi_turn_benchmark(&config).await?;
}
} else if let Some(ref sweep_mc) = cli.sweep_max_concurrency {
// --- Sweep over max-concurrency ---
let values = sweep::parse_concurrency_values(sweep_mc)
.context("Invalid --sweep-max-concurrency")?;
sweep::run_concurrency_sweep(&config, &values, cli.sweep_num_prompts_factor)
.await?;
} else if let Some(ref sweep_rate) = cli.sweep_request_rate {
// --- Sweep over request-rate ---
let values =
sweep::parse_rate_values(sweep_rate).context("Invalid --sweep-request-rate")?;
sweep::run_rate_sweep(&config, &values).await?;
} else if cli.num_runs > 1 {
// --- Multi-run with statistical aggregation ---
multi_run::run_multi(&config, cli.num_runs).await?;
} else {
// --- Normal single benchmark ---
benchmark::run_benchmark(&config).await?;
}
anyhow::Ok(())
})
.context("Benchmark failed")
runtime.block_on(vllm_bench::run(cli.args))
}
rust/src/chat/src/multimodal/item.rs
+4 -4
View File
@@ -38,7 +38,7 @@ pub(super) fn build_batched_items(
let keep_on_cpu = spec.keep_on_cpu_keys.contains(key);
let (value, field) = match spec.field_layout_for(key) {
Some(FieldLayout::Batched) => (
tensor.batched_value_at(index)?,
tensor.batched_wire_value_at(index)?,
MmField::Batched(MmBatchedField { keep_on_cpu }),
),
Some(FieldLayout::Flat { sizes_key }) => {
@@ -47,7 +47,7 @@ pub(super) fn build_batched_items(
})?;
let (start, end) = tensor::flat_range_for_index(sizes, sizes_key, index)?;
(
tensor.flat_value_range(start, end)?,
tensor.flat_wire_value_range(start, end)?,
MmField::Flat(MmFlatField {
slices: vec![MmSlice::Slice(SliceSpec {
start: Some(0),
@@ -60,7 +60,7 @@ pub(super) fn build_batched_items(
)
}
None => (
tensor.clone(),
tensor.try_into()?,
MmField::Shared(MmSharedField {
batch_size: len,
keep_on_cpu,
@@ -71,7 +71,7 @@ pub(super) fn build_batched_items(
data.insert(
key.clone(),
MmFieldElem {
data: Some(value.try_into()?),
data: Some(value),
field,
},
);
rust/src/chat/src/multimodal/tensor.rs
+68 -81
View File
@@ -12,7 +12,7 @@ use vllm_engine_core_client::protocol::tensor::{ShapeExt as _, WireTensor};
use crate::error::{Error, Result, bail_multimodal, multimodal};
/// Representation for multimodal kwarg values for transformation.
#[derive(Debug, Clone)]
#[derive(Debug)]
pub(super) enum KwargValue {
/// Float tensor with row-major flat data and shape.
F32Tensor { data: Vec<f32>, shape: Vec<usize> },
@@ -107,28 +107,19 @@ impl KwargValue {
}
}
impl TryFrom<KwargValue> for ProtocolKwargValue {
impl TryFrom<&KwargValue> for ProtocolKwargValue {
type Error = Error;
fn try_from(value: KwargValue) -> Result<Self> {
match value {
KwargValue::F32Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_f32(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::F16Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_f16(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::Bf16Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_bf16(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::I64Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_i64(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::U32Tensor { data, shape } => Ok(Self::Tensor(
WireTensor::from_u32(shape, data).map_err(Error::Multimodal)?,
)),
KwargValue::Passthrough(value) => Ok(value),
}
fn try_from(value: &KwargValue) -> Result<Self> {
let tensor = match value {
KwargValue::F32Tensor { data, shape } => WireTensor::from_f32(shape.clone(), data),
KwargValue::F16Tensor { data, shape } => WireTensor::from_f16(shape.clone(), data),
KwargValue::Bf16Tensor { data, shape } => WireTensor::from_bf16(shape.clone(), data),
KwargValue::I64Tensor { data, shape } => WireTensor::from_i64(shape.clone(), data),
KwargValue::U32Tensor { data, shape } => WireTensor::from_u32(shape.clone(), data),
KwargValue::Passthrough(value) => return Ok(value.clone()),
};
tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
}
}
@@ -145,63 +136,55 @@ impl KwargValue {
}
}
/// Extract one media item from a batched tensor field.
/// Convert one media item from a batched tensor field to wire bytes.
///
/// Batched fields use their first axis as media-item index and drop that
/// axis in the per-feature value, matching vLLM's batched-field semantics.
pub(super) fn batched_value_at(&self, index: usize) -> Result<Self> {
match self {
Self::F32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::F32Tensor { data, shape })
}
Self::F16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::F16Tensor { data, shape })
}
Self::Bf16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::Bf16Tensor { data, shape })
}
Self::I64Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::I64Tensor { data, shape })
}
Self::U32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, index, index + 1, true)?;
Ok(Self::U32Tensor { data, shape })
}
Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
}
pub(super) fn batched_wire_value_at(&self, index: usize) -> Result<ProtocolKwargValue> {
self.wire_value_range(index, index + 1, true)
}
/// Extract one media item's variable-length range from a flat tensor field.
/// Convert one media item's flat tensor range directly to wire bytes.
///
/// Flat fields keep the first axis as the sliced length for this item.
pub(super) fn flat_value_range(&self, start: usize, end: usize) -> Result<Self> {
match self {
pub(super) fn flat_wire_value_range(
&self,
start: usize,
end: usize,
) -> Result<ProtocolKwargValue> {
self.wire_value_range(start, end, false)
}
fn wire_value_range(
&self,
start: usize,
end: usize,
drop_axis: bool,
) -> Result<ProtocolKwargValue> {
let tensor = match self {
Self::F32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::F32Tensor { data, shape })
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_f32(shape, data)
}
Self::F16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::F16Tensor { data, shape })
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_f16(shape, data)
}
Self::Bf16Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::Bf16Tensor { data, shape })
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_bf16(shape, data)
}
Self::I64Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::I64Tensor { data, shape })
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_i64(shape, data)
}
Self::U32Tensor { data, shape } => {
let (shape, data) = slice_first_axis_range(shape, data, start, end, false)?;
Ok(Self::U32Tensor { data, shape })
let (shape, data) = slice_first_axis_range(shape, data, start, end, drop_axis)?;
WireTensor::from_u32(shape, data)
}
Self::Passthrough(value) => Ok(Self::Passthrough(value.clone())),
}
Self::Passthrough(value) => return Ok(value.clone()),
};
tensor.map(ProtocolKwargValue::Tensor).map_err(Error::Multimodal)
}
}
@@ -240,13 +223,13 @@ fn tensor_as_usize_vec(tensor: &KwargValue) -> Result<Vec<usize>> {
}
/// Slice a flat row-major tensor along its first axis.
fn slice_first_axis_range<T: Clone>(
fn slice_first_axis_range<'a, T>(
shape: &[usize],
data: &[T],
data: &'a [T],
start: usize,
end: usize,
drop_axis: bool,
) -> Result<(Vec<usize>, Vec<T>)> {
) -> Result<(Vec<usize>, &'a [T])> {
let first_dim = *shape.first().ok_or_else(|| multimodal!("tensor has no first dimension"))?;
if start > end || end > first_dim {
bail_multimodal!("invalid tensor slice {start}..{end} for first dimension {first_dim}");
@@ -270,7 +253,7 @@ fn slice_first_axis_range<T: Clone>(
shape[0] = end - start;
shape
};
Ok((out_shape, data[data_start..data_end].to_vec()))
Ok((out_shape, &data[data_start..data_end]))
}
#[cfg(test)]
@@ -278,35 +261,39 @@ mod tests {
use super::*;
#[test]
fn batched_value_at_drops_first_axis() {
fn batched_wire_value_at_drops_first_axis() {
let value = KwargValue::F32Tensor {
data: vec![1.0, 2.0, 3.0, 4.0],
shape: vec![2, 2],
};
let value = value.batched_value_at(1).unwrap();
let ProtocolKwargValue::Tensor(tensor) = value.batched_wire_value_at(1).unwrap() else {
panic!("expected tensor");
};
assert!(matches!(
value,
KwargValue::F32Tensor { data, shape }
if shape == vec![2] && data == vec![3.0, 4.0]
));
assert_eq!(tensor.shape, vec![2]);
assert_eq!(
tensor.data.into_raw_view().unwrap(),
[3.0_f32, 4.0].into_iter().flat_map(f32::to_ne_bytes).collect::<Vec<_>>()
);
}
#[test]
fn flat_value_range_keeps_first_axis() {
fn flat_wire_value_range_keeps_first_axis() {
let value = KwargValue::U32Tensor {
data: (0..10).collect(),
shape: vec![5, 2],
};
let value = value.flat_value_range(1, 3).unwrap();
let ProtocolKwargValue::Tensor(tensor) = value.flat_wire_value_range(1, 3).unwrap() else {
panic!("expected tensor");
};
assert!(matches!(
value,
KwargValue::U32Tensor { data, shape }
if shape == vec![2, 2] && data == vec![2, 3, 4, 5]
));
assert_eq!(tensor.shape, vec![2, 2]);
assert_eq!(
tensor.data.into_raw_view().unwrap(),
[2_u32, 3, 4, 5].into_iter().flat_map(u32::to_ne_bytes).collect::<Vec<_>>()
);
}
#[test]
@@ -336,7 +323,7 @@ mod tests {
let value =
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::BFloat16).unwrap();
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
else {
panic!("expected tensor");
};
@@ -351,7 +338,7 @@ mod tests {
let value =
KwargValue::from_f32_tensor(vec![1.0, -1.0], vec![2], ModelDtype::Float16).unwrap();
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(value).unwrap()
let ProtocolKwargValue::Tensor(tensor) = ProtocolKwargValue::try_from(&value).unwrap()
else {
panic!("expected tensor");
};
rust/src/chat/src/multimodal/video.rs
+4 -4
View File
@@ -130,7 +130,7 @@ fn build_video_item(
let keep_on_cpu = support.spec.keep_on_cpu_keys.contains(&key);
let (value, field) = match support.spec.field_layout_for(&key) {
Some(FieldLayout::Batched) => (
tensor.batched_value_at(0)?,
tensor.batched_wire_value_at(0)?,
MmField::Batched(MmBatchedField { keep_on_cpu }),
),
Some(FieldLayout::Flat { .. }) => {
@@ -138,7 +138,7 @@ fn build_video_item(
.first_dim()
.ok_or_else(|| multimodal!("flat video input `{key}` is not a tensor"))?;
(
tensor,
(&tensor).try_into()?,
MmField::Flat(MmFlatField {
slices: vec![MmSlice::Slice(SliceSpec {
start: Some(0),
@@ -151,7 +151,7 @@ fn build_video_item(
)
}
None => (
tensor,
(&tensor).try_into()?,
MmField::Shared(MmSharedField {
batch_size: 1,
keep_on_cpu,
@@ -162,7 +162,7 @@ fn build_video_item(
data.insert(
key,
MmFieldElem {
data: Some(value.try_into()?),
data: Some(value),
field,
},
);
rust/src/chat/src/renderer/hf/format.rs
+14 -3
View File
@@ -236,9 +236,10 @@ fn has_content_item_loop(root: &Stmt<'_>) -> bool {
loops.into_iter().any(|loop_ast| {
matches!(loop_ast.target, Expr::Var(_))
&& message_varnames
.iter()
.any(|varname| is_var_or_elems_access(&loop_ast.iter, varname, Some("content")))
&& (is_var_access(&loop_ast.iter, "content")
|| message_varnames.iter().any(|varname| {
is_var_or_elems_access(&loop_ast.iter, varname, Some("content"))
}))
})
}
@@ -315,6 +316,16 @@ mod tests {
);
}
#[test]
fn detects_openai_template_with_content_parameter_loop() {
assert_eq!(
detect(
"{% macro render(content) %}{% for item in content %}{{ item }}{% endfor %}{% endmacro %}{% for message in messages %}{{ render(message.content) }}{% endfor %}"
),
ChatTemplateContentFormat::OpenAi
);
}
#[test]
fn detects_openai_template_with_messages_alias() {
assert_eq!(
rust/src/chat/src/renderer/hf/mod.rs
+20
View File
@@ -1309,6 +1309,26 @@ mod tests {
.assert_eq(&rendered);
}
#[test]
fn qwen35_template_auto_detects_openai_multimodal_content() {
let mut request = image_request();
request.chat_options.generation_prompt_mode = GenerationPromptMode::NoGenerationPrompt;
let rendered = render_mm(
QWEN3_5_0_8B_TEMPLATE,
&request,
ChatTemplateContentFormatOption::Auto,
)
.unwrap();
expect![[r#"
Text(
"<|im_start|>user\na<|vision_start|><|image_pad|><|vision_end|>b<|im_end|>\n",
)
"#]]
.assert_debug_eq(&rendered.prompt);
}
#[test]
fn qwen35_template_renders_closed_empty_reasoning_span_when_thinking_disabled() {
let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]);
rust/src/cmd/Cargo.toml
+1
View File
@@ -29,6 +29,7 @@ tokio-util.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
uuid.workspace = true
vllm-bench.workspace = true
vllm-chat.workspace = true
vllm-engine-core-client.workspace = true
vllm-managed-engine.workspace = true
rust/src/cmd/src/cli.rs
+11 -1
View File
@@ -79,13 +79,23 @@ impl Cli {
}
/// Supported top-level CLI commands.
#[derive(Debug, Subcommand, PartialEq, Eq)]
#[derive(Debug, Subcommand)]
pub enum Command {
/// Run the Rust OpenAI frontend as a Python-supervised worker.
Frontend(FrontendArgs),
/// Launch a managed Python headless engine, then run the Rust OpenAI
/// frontend.
Serve(ServeArgs),
/// Run vLLM benchmarks.
#[command(subcommand)]
Bench(BenchCommand),
}
/// Supported benchmark commands.
#[derive(Debug, Subcommand)]
pub enum BenchCommand {
/// Benchmark online serving throughput.
Serve(vllm_bench::BenchServeArgs),
}
/// A JSON-encoded list of strings, matching Python's `json.loads` CLI type for
rust/src/cmd/src/cli/tests.rs
+21 -1
View File
@@ -5,7 +5,27 @@ use expect_test::expect;
use vllm_engine_core_client::TransportMode;
use vllm_server::{Config, HttpListenerMode, ParserSelection, RendererSelection};
use super::{Cli, Command};
use super::{BenchCommand, Cli, Command};
#[test]
fn bench_serve_args_parse_without_managed_engine_repartition() {
let cli = Cli::try_parse_from([
"vllm-rs",
"bench",
"serve",
"--backend",
"openai-chat",
"--request-rate",
"inf",
])
.unwrap();
let Command::Bench(BenchCommand::Serve(args)) = cli.command else {
panic!("expected bench serve args");
};
assert_eq!(args.backend, vllm_bench::BackendKind::OpenaiChat);
assert!(args.request_rate.is_infinite());
}
#[test]
fn serve_args_forward_python_flags_with_separator() {
rust/src/cmd/src/main.rs
+5 -1
View File
@@ -12,7 +12,7 @@ use tokio_util::sync::CancellationToken;
use tracing::{info, warn};
use vllm_managed_engine::ManagedEngineHandle;
use crate::cli::{Cli, Command};
use crate::cli::{BenchCommand, Cli, Command};
#[global_allocator]
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
@@ -100,6 +100,10 @@ fn main() -> Result<()> {
async fn async_main(cli: Cli) -> Result<()> {
match cli.command {
Command::Frontend(args) => vllm_server::serve(args.into_config(), shutdown_signal()).await,
Command::Bench(BenchCommand::Serve(bench_args)) => {
vllm_bench::prepare_process();
vllm_bench::run(bench_args).await
}
Command::Serve(args) => {
let handshake_port = args.managed_engine.resolve_handshake_port()?;
rust/src/engine-core-client/src/protocol/stats.rs
+3
View File
@@ -103,6 +103,9 @@ pub struct PrefillStats {
/// Tokens to be prefilled from external KV transfer.
#[serde(default)]
pub num_external_cached_tokens: u32,
/// Prompt tokens newly admitted into the local prefix cache.
#[serde(default)]
pub num_cache_creation_tokens: u32,
}
/// Stats for debugging the metrics calculation.
rust/src/engine-core-client/src/protocol/tensor.rs
+15 -10
View File
@@ -55,52 +55,57 @@ pub struct WireNdArray {
impl WireNdArray {
/// Build a float32 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_f32(shape: Vec<usize>, data: Vec<f32>) -> Result<Self, String> {
pub fn from_f32(shape: Vec<usize>, data: impl AsRef<[f32]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?;
Ok(Self {
dtype: "float32".to_string(),
shape,
data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(&data)),
data: WireArrayData::RawView(pod_collect_to_vec::<f32, u8>(data)),
})
}
/// Build a float16 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_f16(shape: Vec<usize>, data: Vec<f16>) -> Result<Self, String> {
pub fn from_f16(shape: Vec<usize>, data: impl AsRef<[f16]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?;
Ok(Self {
dtype: "float16".to_string(),
shape,
data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(&data)),
data: WireArrayData::RawView(pod_collect_to_vec::<f16, u8>(data)),
})
}
/// Build a bfloat16 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_bf16(shape: Vec<usize>, data: Vec<bf16>) -> Result<Self, String> {
pub fn from_bf16(shape: Vec<usize>, data: impl AsRef<[bf16]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?;
Ok(Self {
dtype: "bfloat16".to_string(),
shape,
data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(&data)),
data: WireArrayData::RawView(pod_collect_to_vec::<bf16, u8>(data)),
})
}
/// Build an int64 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_i64(shape: Vec<usize>, data: Vec<i64>) -> Result<Self, String> {
pub fn from_i64(shape: Vec<usize>, data: impl AsRef<[i64]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?;
Ok(Self {
dtype: "int64".to_string(),
shape,
data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(&data)),
data: WireArrayData::RawView(pod_collect_to_vec::<i64, u8>(data)),
})
}
/// Build a uint32 tensor/ndarray backed by native-endian raw-view bytes.
pub fn from_u32(shape: Vec<usize>, data: Vec<u32>) -> Result<Self, String> {
pub fn from_u32(shape: Vec<usize>, data: impl AsRef<[u32]>) -> Result<Self, String> {
let data = data.as_ref();
validate_element_count(&shape, data.len())?;
Ok(Self {
dtype: "uint32".to_string(),
shape,
data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(&data)),
data: WireArrayData::RawView(pod_collect_to_vec::<u32, u8>(data)),
})
}
rust/src/llm/src/request_metrics.rs
+1
View File
@@ -375,6 +375,7 @@ mod tests {
num_cached_tokens: 4,
num_local_cached_tokens: 4,
num_external_cached_tokens: 0,
..Default::default()
}),
..Default::default()
},
rust/src/server/src/routes/inference/generate.rs
+56 -7
View File
@@ -238,13 +238,18 @@ fn collect_generate(
None
};
let prompt_logprobs = if include_prompt_logprobs {
let prompt_logprobs = collected.prompt_logprobs.as_ref().ok_or_else(|| {
ApiError::server_error(
"raw generate response requested prompt_logprobs but generation returned none"
.to_string(),
)
})?;
Some(raw_prompt_logprobs_to_maps(prompt_logprobs))
match collected.prompt_logprobs.as_ref() {
Some(prompt_logprobs) => Some(raw_prompt_logprobs_to_maps(prompt_logprobs)),
// A single-token prompt has no scored positions; same mapping
// as /v1/completions.
None if collected.prompt_token_ids.len() == 1 => Some(vec![None]),
None => {
return Err(ApiError::server_error(
"raw generate response requested prompt_logprobs but generation returned none"
.to_string(),
));
}
}
} else {
None
};
@@ -472,4 +477,48 @@ mod tests {
Some(2)
);
}
#[test]
fn collect_generate_maps_prompt_logprobs_for_single_token_prompt() {
let output_without_payload = |prompt_token_ids: Vec<u32>| CollectedGenerateOutput {
request_id: "raw-1".to_string(),
prompt_logprobs: None,
token_ids: vec![3],
logprobs: None,
finish_reason: FinishReason::stop_eos(),
usage: vllm_llm::TokenUsage {
prompt_token_count: prompt_token_ids.len(),
output_token_count: 1,
cached_token_count: 0,
},
kv_transfer_params: None,
ec_transfer_params: None,
prompt_token_ids,
};
let response = collect_generate(
output_without_payload(vec![9707]),
"raw-1".to_string(),
ApiServerOptions::default(),
ResponseOptions {
include_prompt_logprobs: true,
..Default::default()
},
)
.expect("single-token prompt without payload maps to [None]");
let prompt_logprobs = response.prompt_logprobs.expect("prompt logprobs present");
assert_eq!(prompt_logprobs.len(), 1);
assert!(prompt_logprobs[0].is_none());
collect_generate(
output_without_payload(vec![9707, 11]),
"raw-2".to_string(),
ApiServerOptions::default(),
ResponseOptions {
include_prompt_logprobs: true,
..Default::default()
},
)
.expect_err("multi-token prompt without payload is an engine failure");
}
}
rust/src/server/src/routes/openai/chat_completions.rs
+5 -8
View File
@@ -35,7 +35,7 @@ use crate::routes::openai::chat_completions::types::{
ChatMessageDelta,
};
use crate::routes::openai::utils::logprobs::{
decoded_logprobs_to_openai_chat, decoded_prompt_logprobs_to_maps,
decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps,
};
use crate::routes::openai::utils::types::{
ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
@@ -181,14 +181,11 @@ async fn collect_chat_completion(
None
};
let prompt_logprobs = if include_prompt_logprobs {
Some(decoded_prompt_logprobs_to_maps(
prompt_logprobs.as_ref().ok_or_else(|| {
server_error!(
"chat response requested prompt_logprobs but generation returned none"
)
})?,
Some(prompt_logprobs_to_maps(
prompt_logprobs.as_ref(),
&prompt_token_ids,
return_tokens_as_token_ids,
))
)?)
} else {
None
};
rust/src/server/src/routes/openai/completions.rs
+2 -24
View File
@@ -5,7 +5,6 @@ mod convert;
mod types;
mod validate;
use std::collections::HashMap;
use std::convert::Infallible;
use std::result::Result;
use std::sync::Arc;
@@ -29,8 +28,8 @@ use vllm_text::{
use self::convert::{ResponseOptions, prepare_completion_request};
use super::utils::logprobs::{
collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_maps,
decoded_prompt_logprobs_to_openai, text_len,
collected_logprobs_to_openai, decoded_logprobs_to_openai, decoded_prompt_logprobs_to_openai,
prompt_logprobs_to_maps, text_len,
};
use super::utils::types::Usage;
use crate::config::ApiServerOptions;
@@ -505,27 +504,6 @@ fn prompt_only_logprobs_to_openai(
))
}
fn prompt_logprobs_to_maps(
prompt_logprobs: Option<&DecodedPromptLogprobs>,
prompt_token_ids: &[u32],
return_tokens_as_token_ids: bool,
) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
if let Some(prompt_logprobs) = prompt_logprobs {
return Ok(decoded_prompt_logprobs_to_maps(
prompt_logprobs,
return_tokens_as_token_ids,
));
}
if let [_token_id] = prompt_token_ids {
return Ok(vec![None]);
}
Err(server_error!(
"completion response requested prompt_logprobs but generation returned none"
))
}
fn usage_chunk(
request_id: &str,
response_model: &str,
rust/src/server/src/routes/openai/utils/logprobs.rs
+31 -14
View File
@@ -100,20 +100,31 @@ pub fn decoded_prompt_logprobs_to_openai(
})
}
/// Convert decoded prompt logprobs into the vLLM-style prompt-logprobs response
/// shape.
pub fn decoded_prompt_logprobs_to_maps(
prompt_logprobs: &DecodedPromptLogprobs,
/// Map decoded prompt logprobs into vLLM-style per-position maps, treating a
/// missing single-token payload as `[None]`.
pub fn prompt_logprobs_to_maps(
prompt_logprobs: Option<&DecodedPromptLogprobs>,
prompt_token_ids: &[u32],
return_tokens_as_token_ids: bool,
) -> Vec<Option<HashMap<String, f32>>> {
std::iter::once(None)
.chain(prompt_logprobs.scored_positions.iter().map(|position| {
Some(position_top_logprobs_map(
position,
return_tokens_as_token_ids,
))
}))
.collect()
) -> Result<Vec<Option<HashMap<String, f32>>>, ApiError> {
if let Some(prompt_logprobs) = prompt_logprobs {
return Ok(std::iter::once(None)
.chain(prompt_logprobs.scored_positions.iter().map(|position| {
Some(position_top_logprobs_map(
position,
return_tokens_as_token_ids,
))
}))
.collect());
}
if let [_token_id] = prompt_token_ids {
return Ok(vec![None]);
}
Err(server_error!(
"prompt_logprobs were requested but generation returned none"
))
}
/// Convert decoded token-position logprobs into the OpenAI chat `logprobs`
@@ -275,7 +286,13 @@ pub fn clamp_logprob(logprob: f32) -> f32 {
mod tests {
use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
use super::decoded_logprobs_to_openai_chat;
use super::{decoded_logprobs_to_openai_chat, prompt_logprobs_to_maps};
#[test]
fn prompt_logprobs_maps_reject_missing_multi_token_payload() {
prompt_logprobs_to_maps(None, &[9707, 11], false)
.expect_err("multi-token prompt without payload is an engine failure");
}
fn sample_logprobs() -> DecodedLogprobs {
DecodedLogprobs {
setup.py
+1 -1
View File
@@ -1268,7 +1268,7 @@ setup(
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
"tensorizer": ["tensorizer==2.10.1"],
"fastsafetensors": ["fastsafetensors >= 0.3.2"],
"instanttensor": ["instanttensor >= 0.1.5"],
"instanttensor": ["instanttensor >= 0.1.9"],
"runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
"audio": [
"av",
tests/config/test_model_arch_config.py
+17
View File
@@ -131,6 +131,23 @@ def test_head_size_falls_back_when_head_dim_is_zero():
assert convertor.get_head_size() == 128
def test_legacy_modelopt_config_without_producer_is_normalized():
quantization_config = {
"quantization": {
"quant_algo": "NVFP4",
"group_size": 16,
"kv_cache_quant_algo": None,
"exclude_modules": [],
"modelopt_quant_config": {"quant_cfg": {}},
}
}
hf_config = PretrainedConfig(quantization_config=quantization_config)
convertor = ModelArchConfigConvertorBase(hf_config, hf_config)
assert convertor.get_quantization_config()["quant_method"] == "modelopt_fp4"
@pytest.mark.parametrize("model", BASE_MODELS_TO_TEST)
def test_base_model_arch_config(model: str):
"""Test model architecture config for base models."""
tests/entrypoints/anthropic/test_anthropic_messages_conversion.py
+25 -60
View File
@@ -23,7 +23,6 @@ from vllm.entrypoints.anthropic.protocol import (
from vllm.entrypoints.anthropic.serving import (
AnthropicServingMessages,
_build_anthropic_usage,
_get_cached_tokens,
)
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionResponse,
@@ -668,42 +667,6 @@ class TestThinkingBlockConversion:
# ======================================================================
class TestGetCachedTokens:
"""Tests for _get_cached_tokens helper."""
def test_none_usage(self):
assert _get_cached_tokens(None) is None
def test_no_prompt_tokens_details(self):
usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
assert _get_cached_tokens(usage) is None
def test_cached_tokens_present(self):
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
)
assert _get_cached_tokens(usage) == 80
def test_cached_tokens_zero(self):
"""Zero cached tokens should return 0, not None."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
)
assert _get_cached_tokens(usage) == 0
def test_cached_tokens_none_in_details(self):
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=None),
)
assert _get_cached_tokens(usage) is None
class TestBuildAnthropicUsage:
"""Tests for _build_anthropic_usage helper.
@@ -711,36 +674,32 @@ class TestBuildAnthropicUsage:
vLLM's prompt_tokens is the total.
"""
def test_no_cache_info(self):
"""When cache info is unavailable, return raw prompt_tokens."""
result = _build_anthropic_usage(100, 10, None)
assert result.input_tokens == 100
assert result.output_tokens == 10
assert result.cache_read_input_tokens is None
assert result.cache_creation_input_tokens is None
def test_cache_hit(self):
"""When cache is hit, input_tokens excludes cached tokens."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
prompt_tokens_details=PromptTokenUsageInfo(
cached_tokens=80, created_cache_tokens=10
),
)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 20 # 100 - 80
result = _build_anthropic_usage(usage)
assert result.input_tokens == 10 # 100 - 80 - 10
assert result.output_tokens == 10
assert result.cache_read_input_tokens == 80
assert result.cache_creation_input_tokens == 0
assert result.cache_creation_input_tokens == 10
def test_zero_cached_tokens(self):
"""Zero cached tokens should still set cache_creation to 0."""
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
prompt_tokens_details=PromptTokenUsageInfo(
cached_tokens=0, created_cache_tokens=0
),
)
result = _build_anthropic_usage(100, 10, usage)
assert result.input_tokens == 100 # 100 - 0
result = _build_anthropic_usage(usage)
assert result.input_tokens == 100 # 100 - 0 - 0
assert result.cache_read_input_tokens == 0
assert result.cache_creation_input_tokens == 0
@@ -749,9 +708,11 @@ class TestBuildAnthropicUsage:
usage = UsageInfo(
prompt_tokens=100,
completion_tokens=10,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=100),
prompt_tokens_details=PromptTokenUsageInfo(
cached_tokens=100, created_cache_tokens=0
),
)
result = _build_anthropic_usage(100, 10, usage)
result = _build_anthropic_usage(usage)
assert result.input_tokens == 0
assert result.cache_read_input_tokens == 100
assert result.cache_creation_input_tokens == 0
@@ -759,7 +720,7 @@ class TestBuildAnthropicUsage:
def test_no_prompt_tokens_details(self):
"""UsageInfo without prompt_tokens_details returns no cache info."""
usage = UsageInfo(prompt_tokens=100, completion_tokens=10)
result = _build_anthropic_usage(100, 10, usage)
result = _build_anthropic_usage(usage)
assert result.input_tokens == 100
assert result.cache_read_input_tokens is None
assert result.cache_creation_input_tokens is None
@@ -1241,7 +1202,9 @@ class TestStreamingCacheUsageSemantics:
prompt_tokens=100,
completion_tokens=5,
total_tokens=105,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=80),
prompt_tokens_details=PromptTokenUsageInfo(
cached_tokens=80, created_cache_tokens=10
),
),
)
yield "data: [DONE]"
@@ -1263,9 +1226,9 @@ class TestStreamingCacheUsageSemantics:
delta_usage = next(
data["usage"] for ev, data in events if ev == "message_delta"
)
assert delta_usage["input_tokens"] == 20 # 100 - 80
assert delta_usage["input_tokens"] == 10 # 100 - 80 - 10
assert delta_usage["cache_read_input_tokens"] == 80
assert delta_usage["cache_creation_input_tokens"] == 0
assert delta_usage["cache_creation_input_tokens"] == 10
@pytest.mark.asyncio
async def test_streaming_no_cache_hit(self):
@@ -1284,7 +1247,9 @@ class TestStreamingCacheUsageSemantics:
prompt_tokens=50,
completion_tokens=5,
total_tokens=55,
prompt_tokens_details=PromptTokenUsageInfo(cached_tokens=0),
prompt_tokens_details=PromptTokenUsageInfo(
cached_tokens=0, created_cache_tokens=0
),
),
)
yield "data: [DONE]"
@@ -1302,7 +1267,7 @@ class TestStreamingCacheUsageSemantics:
assert start_usage["input_tokens"] == 50
assert "cache_read_input_tokens" not in start_usage
assert "cache_creation_input_tokens" not in start_usage
assert delta_usage["input_tokens"] == 50 # 50 - 0
assert delta_usage["input_tokens"] == 50 # 50 - 0 - 0
assert delta_usage["cache_read_input_tokens"] == 0
assert delta_usage["cache_creation_input_tokens"] == 0
tests/entrypoints/anthropic/test_messages.py
+52
View File
@@ -18,6 +18,7 @@ def server():
"--max-model-len",
"2048",
"--enforce-eager",
"--enable-prompt-tokens-details",
"--enable-auto-tool-choice",
"--tool-call-parser",
"hermes",
@@ -191,3 +192,54 @@ async def test_anthropic_structured_output(client: anthropic.AsyncAnthropic):
json_obj = json.loads(response.content[0].text)
for key in ["name", "email", "plan_interest", "demo_requested"]:
assert key in json_obj, f"Missing key in output: {key}"
@pytest.mark.asyncio
async def test_anthropic_streaming_cache_usage(client: anthropic.AsyncAnthropic):
async def get_stream_usage(resp):
prompt_tokens = None
usage = None
async for chunk in resp:
if (
chunk.type == "message_start"
and chunk.message is not None
and chunk.message.usage is not None
):
prompt_tokens = chunk.message.usage.input_tokens
elif chunk.type == "message_delta" and chunk.usage is not None:
usage = chunk.usage
assert usage is not None
assert usage.input_tokens >= 0
assert usage.output_tokens >= 0
cache_created = usage.cache_creation_input_tokens
cache_read = usage.cache_read_input_tokens
assert cache_read is not None
assert cache_created is not None
assert cache_created >= 0
assert cache_read >= 0
assert prompt_tokens == usage.input_tokens + cache_created + cache_read
return usage
request = dict(
model="claude-3-7-sonnet-latest",
max_tokens=1,
temperature=0.0,
messages=[
{
"role": "user",
"content": "Cache coverage sentinel. " * 256
+ "Answer with exactly one word: ok.",
}
],
stream=True,
)
cold_usage = await get_stream_usage(await client.messages.create(**request))
assert cold_usage.cache_read_input_tokens == 0
assert cold_usage.cache_creation_input_tokens is not None
assert cold_usage.cache_creation_input_tokens > 0
warm_usage = await get_stream_usage(await client.messages.create(**request))
assert warm_usage.cache_read_input_tokens is not None
assert warm_usage.cache_read_input_tokens > 0
tests/entrypoints/openai/chat_completion/test_chat_error.py
+12
View File
@@ -515,3 +515,15 @@ def test_structured_outputs_structural_tag_invalid(structural_tag):
messages=[{"role": "user", "content": "hello"}],
structured_outputs={"structural_tag": structural_tag},
)
@pytest.mark.parametrize("field_name", ["prompt_logprobs", "top_logprobs"])
def test_non_numeric_logprobs_rejected(field_name):
"""A non-numeric logprobs value must be a clean 400 validation error, not a
TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
ChatCompletionRequest(
model=MODEL_NAME,
messages=[{"role": "user", "content": "hello"}],
**{field_name: "2"},
)
tests/entrypoints/openai/chat_completion/test_serving_chat.py
+12 -5
View File
@@ -831,16 +831,23 @@ def test_mm_prompt_tokens_details():
assert counts == {"image": 600, "video": 1200}
# Gated off, or nothing to report -> no details.
assert _make_prompt_tokens_details(False, 5, counts) is None
assert _make_prompt_tokens_details(True, None, None) is None
assert _make_prompt_tokens_details(False, 5, 0, counts) is None
assert _make_prompt_tokens_details(True, None, None, None) is None
# Zero cached_tokens is still reported (not None), matching the cached-only
# behavior; multimodal counts ride alongside even when cached_tokens is None.
assert _make_prompt_tokens_details(True, 0, None).cached_tokens == 0
details = _make_prompt_tokens_details(True, None, counts)
details = _make_prompt_tokens_details(True, 0, 0, None)
assert details.cached_tokens == 0
assert details.created_cache_tokens == 0
assert details.multimodal_tokens is None
details = _make_prompt_tokens_details(True, None, None, counts)
assert details.cached_tokens is None
assert details.created_cache_tokens is None
assert details.multimodal_tokens == {"image": 600, "video": 1200}
details = _make_prompt_tokens_details(True, 3, 0, counts)
assert details.cached_tokens == 3
assert details.created_cache_tokens == 0
assert details.multimodal_tokens == {"image": 600, "video": 1200}
assert _make_prompt_tokens_details(True, 3, counts).cached_tokens == 3
@pytest.mark.asyncio
tests/entrypoints/openai/completion/test_completion_error.py
+13
View File
@@ -610,3 +610,16 @@ class TestCompletionPromptListLimit:
max_tokens=1,
)
assert len(request.prompt_embeds) == 5
@pytest.mark.parametrize("field_name", ["prompt_logprobs", "logprobs"])
def test_non_numeric_logprobs_rejected(field_name):
"""A non-numeric logprobs value must be a clean 400 validation error, not a
TypeError from the mode='before' comparison (which surfaces as HTTP 500)."""
with pytest.raises(ValidationError, match=f"`{field_name}` must be an integer"):
CompletionRequest(
model=MODEL_NAME,
prompt="Test prompt",
max_tokens=10,
**{field_name: "2"},
)
tests/entrypoints/openai/responses/test_parsable_context.py
+1 -1
View File
@@ -68,7 +68,7 @@ async def client(server):
async def test_basic(client: OpenAI, model_name: str):
response = await client.responses.create(
model=model_name,
input="What is 123 * 456?",
input="What is 123 * 456? Answer with only the number.",
temperature=0.0,
)
assert response is not None
tests/entrypoints/openai/responses/test_sampling_params.py
+15
View File
@@ -132,6 +132,21 @@ class TestResponsesRequestSamplingParams:
assert sampling_params.structured_outputs is not None
assert sampling_params.structured_outputs.grammar == "root ::= 'hello'"
def test_text_format_json_object_enables_structured_outputs(self):
"""text.format json_object enables structured outputs for sampling."""
request = ResponsesRequest(
model="test-model",
input="test input",
text=ResponseTextConfig.model_validate({"format": {"type": "json_object"}}),
)
sampling_params = request.to_sampling_params(default_max_tokens=1000)
assert sampling_params.structured_outputs is not None
assert sampling_params.structured_outputs.json_object is True
assert sampling_params.structured_outputs.json is None
assert request.structured_outputs is None
def test_structured_outputs_and_json_schema_conflict(self):
"""Test that specifying both structured_outputs and json_schema raises."""
structured_outputs = StructuredOutputsParams(grammar="root ::= 'hello'")
tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
+16
View File
@@ -0,0 +1,16 @@
model_name: "nvidia/GLM-5.2-NVFP4"
accuracy_threshold: 0.90
num_questions: 1319
num_fewshot: 5
max_concurrency: 100
server_args: >-
--enforce-eager
--max-model-len 4096
--safetensors-load-strategy prefetch
--moe-backend flashinfer_cutlass
--prefill-context-parallel-size 4
--enable-expert-parallel
--kv-cache-dtype fp8
env:
VLLM_LOGGING_LEVEL: "DEBUG"
VLLM_USE_V2_MODEL_RUNNER: "1"
tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
+17
View File
@@ -0,0 +1,17 @@
model_name: "nvidia/GLM-5.2-NVFP4"
accuracy_threshold: 0.90
num_questions: 1319
num_fewshot: 5
max_concurrency: 100
server_args: >-
--enforce-eager
--max-model-len 4096
--safetensors-load-strategy prefetch
--moe-backend flashinfer_cutlass
--tensor-parallel-size 2
--prefill-context-parallel-size 2
--enable-expert-parallel
--kv-cache-dtype fp8
env:
VLLM_LOGGING_LEVEL: "DEBUG"
VLLM_USE_V2_MODEL_RUNNER: "1"
tests/evals/gsm8k/configs/models-pcp.txt
+2
View File
@@ -0,0 +1,2 @@
GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
tests/evals/gsm8k/gsm8k_eval.py
+15 -1
View File
@@ -217,6 +217,7 @@ def evaluate_gsm8k(
seed: int | None = 42,
request_timeout_seconds: float = 600,
gen_prefix: str = "",
max_concurrency: int | None = None,
) -> dict[str, float | int]:
"""
Evaluate GSM8K accuracy using vLLM serve endpoint.
@@ -261,7 +262,14 @@ def evaluate_gsm8k(
return answer, tokens
timeout = aiohttp.ClientTimeout(total=request_timeout_seconds)
async with aiohttp.ClientSession(timeout=timeout) as session:
connector = (
aiohttp.TCPConnector(limit=max_concurrency)
if max_concurrency is not None
else None
)
async with aiohttp.ClientSession(
timeout=timeout, connector=connector
) as session:
tasks = [get_answer(session, i) for i in range(num_questions)]
await tqdm.gather(*tasks, desc="Evaluating")
@@ -343,6 +351,11 @@ def main() -> None:
parser.add_argument(
"--seed", type=int, default=42, help="Random seed for reproducibility"
)
parser.add_argument(
"--max-concurrency",
type=int,
help="Maximum number of concurrent requests",
)
parser.add_argument("--save-results", type=str, help="Save results to JSON file")
args = parser.parse_args()
@@ -355,6 +368,7 @@ def main() -> None:
port=args.port,
temperature=args.temperature,
seed=args.seed,
max_concurrency=args.max_concurrency,
)
# Print results to terminal
tests/evals/gsm8k/test_gsm8k_correctness.py
+1
View File
@@ -59,6 +59,7 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict:
seed=eval_config.get("seed", 42),
request_timeout_seconds=request_timeout_seconds,
gen_prefix=eval_config.get("gen_prefix", ""),
max_concurrency=eval_config.get("max_concurrency"),
)
return results
tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py
+2
View File
@@ -58,6 +58,7 @@ def _make_builder():
max_num_batched_tokens + 1, dtype=torch.int32, device="cpu"
)
builder._num_attention_heads = 16
builder._num_compute_units = current_platform.num_compute_units()
builder._mla_work_meta_data = torch.empty(1, dtype=torch.int32, device="cpu")
builder._mla_work_indptr = torch.empty(1, dtype=torch.int32, device="cpu")
builder._mla_work_info_set = torch.empty(1, dtype=torch.int32, device="cpu")
@@ -116,6 +117,7 @@ def test_sparse_persistent_metadata_syncs_only_after_recompute(monkeypatch):
assert events == ["metadata", "sync"]
assert fake_get_mla_metadata_v1_mock.call_count == 1
assert fake_get_mla_metadata_v1_mock.call_args.kwargs["max_split_per_batch"] == 1
events.clear()
tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+1 -1
View File
@@ -425,7 +425,7 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No
match the single-call result.
"""
from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
causal_conv1d_torch,
causal_conv1d_fn_cpu as causal_conv1d_torch,
)
x, weight, bias = _conv_inputs(total_tokens)
tests/kernels/mamba/test_causal_conv1d.py
+8 -3
View File
@@ -18,8 +18,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
DEVICE = current_platform.device_type
pytestmark = pytest.mark.skipif(
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
reason="causal_conv1d Triton kernels require CUDA-alike or XPU",
not (
current_platform.is_cuda_alike()
or current_platform.is_xpu()
or current_platform.is_cpu()
),
reason="causal_conv1d Triton kernels require CUDA-alike, XPU, or CPU",
)
@@ -284,7 +288,8 @@ def test_causal_conv1d_varlen(
batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype
):
device = DEVICE
torch.accelerator.empty_cache()
if not current_platform.is_cpu():
torch.accelerator.empty_cache()
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
if itype == torch.bfloat16:
rtol, atol = 1e-2, 5e-2
tests/kernels/mamba/test_mamba_ssm.py
+53 -3
View File
@@ -20,8 +20,12 @@ from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
DEVICE = current_platform.device_type
pytestmark = pytest.mark.skipif(
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
reason="mamba_ssm kernels require CUDA-alike or XPU",
not (
current_platform.is_cuda_alike()
or current_platform.is_xpu()
or current_platform.is_cpu()
),
reason="mamba_ssm kernels require CUDA-alike, XPU, or CPU",
)
# selective_scan_fn is backed by the CUDA-only `ops.selective_scan_fwd` C++ op,
@@ -342,12 +346,23 @@ def test_selective_scan(
@pytest.mark.parametrize("has_z", [False, True])
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update(dim, dstate, has_z, itype):
device = DEVICE
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
if itype == torch.bfloat16:
rtol, atol = 1e-2, 5e-2
if current_platform.is_rocm() or current_platform.is_xpu():
if (
current_platform.is_rocm()
or current_platform.is_xpu()
or current_platform.is_device_capability_family(90)
):
atol *= 2
# set seed
set_random_seed(0)
@@ -432,6 +447,13 @@ def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_r
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
@pytest.mark.parametrize("max_seq_len", [1, 2, 4])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):
device = DEVICE
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)
@@ -693,6 +715,13 @@ def test_selective_scan_varlen(
@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
# tests correctness in case subset of the sequences are padded
@pytest.mark.parametrize("with_padding", [True, False])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update_with_batch_indices(
with_padding, dim, dstate, has_z, itype
):
@@ -785,6 +814,13 @@ def test_selective_state_update_with_batch_indices(
@pytest.mark.parametrize("ngroups", [1, 4])
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 4096])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update_with_heads_with_batch_indices(
dim, dstate, ngroups, has_z, tie_hdim, itype
):
@@ -858,6 +894,13 @@ def test_selective_state_update_with_heads_with_batch_indices(
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 4096])
@pytest.mark.parametrize("max_seq_len", [2, 4])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update_with_num_accepted_tokens(
dim, dstate, has_z, itype, max_seq_len
):
@@ -984,6 +1027,13 @@ def test_selective_state_update_with_num_accepted_tokens(
@pytest.mark.parametrize("dstate", [16, 64])
@pytest.mark.parametrize("dim", [2048, 4096])
@pytest.mark.parametrize("max_seq_len", [2, 4])
@pytest.mark.skipif(
current_platform.is_cpu(),
reason=(
"CPU kernel for selective_state_update only supports "
"Mamba 2 (scalar A/dt), not Mamba 1."
),
)
def test_selective_state_update_varlen_with_num_accepted(
dim, dstate, has_z, itype, max_seq_len
):
tests/kernels/moe/test_ocp_mx_moe.py
+46
View File
@@ -1512,3 +1512,49 @@ def test_rocm_mxfp4_moe_oracle(
# Check accuracy using per-backend thresholds
check_accuracy(ref, out, atol=0.1, rtol=config["rtol"], percent=config["percent"])
# -----------------------------------------------------------------------------
# MXFP4 emulation size-rounding tests
# -----------------------------------------------------------------------------
# Emulation needs each per-partition dim rounded up to OCP_MX_BLOCK_SIZE (32);
# a non-block-aligned shard (e.g. GPT-OSS 2880 // 4 = 720) otherwise truncates
# the scale buffer and fails weight loading.
# NOTE: gated to ROCm since it is the emulation backend's current target;
# remove this skip if the backend is enabled on non-ROCm platforms.
@pytest.mark.skipif(not ROCM_AVAILABLE, reason="emulation backend targets ROCm")
@pytest.mark.parametrize(
"hidden_size,intermediate_size,expected_hidden,expected_intermediate",
[
(2880, 720, 2880, 736), # GPT-OSS TP=4 shard: 720 -> round_up(720, 32)
(2880, 360, 2880, 384), # GPT-OSS TP=8 shard: 360 -> 384
(2880, 2880, 2880, 2880), # already block-aligned: unchanged
(90, 90, 96, 96), # both dims unaligned
],
)
def test_mxfp4_emulation_rounds_up_to_block_size(
hidden_size: int,
intermediate_size: int,
expected_hidden: int,
expected_intermediate: int,
):
"""Emulation must block-align per-partition dims to OCP_MX_BLOCK_SIZE."""
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
Mxfp4MoeBackend,
mxfp4_round_up_hidden_size_and_intermediate_size,
)
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
OCP_MX_BLOCK_SIZE,
)
rounded_hidden, rounded_intermediate = (
mxfp4_round_up_hidden_size_and_intermediate_size(
Mxfp4MoeBackend.EMULATION, hidden_size, intermediate_size
)
)
assert rounded_hidden == expected_hidden
assert rounded_intermediate == expected_intermediate
# The block-scale buffer (dim // OCP_MX_BLOCK_SIZE) must not floor-truncate.
assert rounded_hidden % OCP_MX_BLOCK_SIZE == 0
assert rounded_intermediate % OCP_MX_BLOCK_SIZE == 0
tests/kernels/moe/test_topk_softplus_sqrt.py
+45
View File
@@ -10,6 +10,7 @@ from vllm.model_executor.layers.fused_moe.config import (
RoutingMethodType,
get_routing_method_type,
)
from vllm.model_executor.layers.fused_moe.router.dsv4_topk import dsv4_topk
from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
fused_topk_bias,
)
@@ -186,3 +187,47 @@ def test_fused_topk_softplus_sqrt_hash(
sorted_w_ref = topk_weights_ref.gather(1, idx_ref)
sorted_w = topk_weights.gather(1, idx_ops)
torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)
@pytest.mark.skipif(
not current_platform.is_cuda(),
reason="The DeepSeek V4 fast path is CUDA-only.",
)
@pytest.mark.parametrize(
("num_tokens", "num_experts", "indices_type"),
[
(0, 256, torch.uint32),
(17, 256, torch.uint32),
(17, 384, torch.int64),
],
)
def test_dsv4_fast_topk(
num_tokens: int,
num_experts: int,
indices_type: torch.dtype,
):
torch.manual_seed(0)
gating_output = torch.randn(
(num_tokens, num_experts), dtype=torch.float32, device="cuda"
)
correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda")
topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
gating_output=gating_output,
topk=6,
renormalize=True,
routed_scaling_factor=1.5,
e_score_correction_bias=correction_bias,
)
topk_weights, topk_ids = dsv4_topk(
gating_output, correction_bias, indices_type, 1.5
)
assert topk_ids.dtype == indices_type
torch.testing.assert_close(topk_ids_ref.to(indices_type), topk_ids, atol=0, rtol=0)
torch.testing.assert_close(
topk_weights_ref,
topk_weights,
atol=2e-5,
rtol=2e-5,
)
tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py
-2
View File
@@ -33,8 +33,6 @@ def test_instanttensor_model_loader():
hf_safetensors_tensors = {}
for name, tensor in instanttensor_weights_iterator(safetensors, True):
# Copy the tensor immediately as it is a reference to the internal
# buffer of instanttensor.
instanttensor_tensors[name] = tensor.to("cpu")
for name, tensor in safetensors_weights_iterator(safetensors, True):
tests/models/inkling/test_moe_weight_layout.py
+25
View File
@@ -7,6 +7,7 @@ import pytest
import torch
from vllm.lora.utils import get_supported_lora_modules
from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config
from vllm.models.inkling.nvidia import moe
from vllm.models.inkling.nvidia.model import _TmlForCausalLMBase
from vllm.platforms import current_platform
@@ -87,6 +88,30 @@ def test_custom_embedding_is_not_a_lora_target() -> None:
assert "lm_head" in supported
def test_inkling_mapper_maps_modelopt_exclusions() -> None:
quant_config = ModelOptNvFp4Config.from_config(
{
"quantization": {
"quant_algo": "NVFP4",
"group_size": 16,
"kv_cache_quant_algo": None,
"exclude_modules": [
"model.llm.layers.2.mlp.experts",
"model.llm.layers.2.mlp.shared_experts",
],
}
}
)
quant_config.apply_vllm_mapper(
_TmlForCausalLMBase.hf_to_vllm_mapper.get_unstacked_mapper()
)
assert quant_config.is_layer_excluded("model.layers.2.mlp.experts")
assert quant_config.is_layer_excluded("model.layers.2.mlp.shared_experts")
assert not quant_config.is_layer_excluded("model.layers.3.mlp.experts")
@pytest.mark.parametrize(("projection", "amax"), [("w13", 4.375), ("w2", 2960.0)])
def test_moe_loads_calibrated_input_scale(projection: str, amax: float) -> None:
experts = SimpleNamespace(
tests/models/language/generation/test_common.py
+11
View File
@@ -188,6 +188,16 @@ def test_models(
prompt_embeds.append(embed.squeeze(0))
vllm_kwargs = {}
if (
model == "bigscience/bloom-560m"
and current_platform.is_device_capability_family(90)
):
# On SM90, the metadata builder otherwise selects FA3 AOT scheduling
# before Bloom's ALiBi layers fall back to FA2. Pinning FA2 keeps the
# builder and layer consistent and preserves the L4 test path.
vllm_kwargs["attention_config"] = {"flash_attn_version": 2}
with vllm_runner(
model,
tokenizer_name=model_info.tokenizer or model,
@@ -200,6 +210,7 @@ def test_models(
max_num_seqs=1 if current_platform.is_rocm() else 2,
enable_prompt_embeds=use_prompt_embeds,
compilation_config={"cudagraph_capture_sizes": [1, 2]},
**vllm_kwargs,
) as vllm_model:
vllm_outputs = vllm_model.generate_greedy_logprobs(
example_prompts, max_tokens, num_logprobs
tests/models/language/generation/test_hybrid.py
+6 -1
View File
@@ -384,8 +384,13 @@ def test_fp32_cache_state(
example_prompts, max_tokens, num_logprobs
)
# Leave enough headroom for repeated engine initialization on a
# 32.5 GiB MIG.
with vllm_runner(
model, max_num_seqs=MAX_NUM_SEQS, **{cache_dtype_param: "float32"}
model,
max_num_seqs=MAX_NUM_SEQS,
gpu_memory_utilization=0.9,
**{cache_dtype_param: "float32"},
) as vllm_model:
vllm_outputs = vllm_model.generate_greedy_logprobs(
example_prompts, max_tokens, num_logprobs
tests/models/multimodal/processing/test_glm4_1v.py
+52
View File
@@ -1,9 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import Mock
import pytest
from vllm.assets.video import VideoAsset
from vllm.model_executor.models.glm4_1v import (
Glm4vForConditionalGeneration,
Glm4vProcessingInfo,
)
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import batched_tensors_equal
from vllm.multimodal.video import DynamicVideoBackend, VideoBackend
@@ -11,6 +17,52 @@ from vllm.multimodal.video import DynamicVideoBackend, VideoBackend
from ...utils import build_model_context
@pytest.mark.parametrize(
(
"max_video_pixels",
"max_tokens",
"expected_num_frames",
),
[
(47_040_000, 124_988, 11),
(47_040_000, 30_000, 24),
(100_352_000, 124_988, 21),
(100_352_000, 30_000, 7),
(100_352_000, 0, 1),
],
)
def test_get_max_video_frames_matches_glm_resize(
max_video_pixels: int,
max_tokens: int,
expected_num_frames: int,
):
info = Mock(spec=Glm4vProcessingInfo)
info.get_image_size_with_most_features.return_value = (2184, 2184)
info._get_video_max_pixels.return_value = max_video_pixels
vision_config = info.get_hf_config.return_value.vision_config
vision_config.patch_size = 14
vision_config.spatial_merge_size = 2
vision_config.temporal_patch_size = 2
info._get_vision_info.side_effect = lambda **kwargs: (
Glm4vProcessingInfo._get_vision_info(info, **kwargs)
)
num_frames = Glm4vProcessingInfo._get_max_video_frames(
info,
max_tokens=max_tokens,
)
assert num_frames == expected_num_frames
assert info._get_video_max_pixels.call_count == 1
assert info._get_vision_info.call_count == 600
def test_encoder_cudagraph_uses_model_video_frame_limit():
model = Mock()
assert Glm4vForConditionalGeneration.get_max_frames_per_video(model) == 600
@pytest.mark.parametrize("model_id", ["zai-org/GLM-4.1V-9B-Thinking"])
@pytest.mark.parametrize("expected_toks_per_frame", [299])
@pytest.mark.parametrize(
tests/models/quantization/test_fp8.py
+2
View File
@@ -67,6 +67,8 @@ def test_models(
if kv_cache_dtype == "fp8_e5m2" and current_platform.is_rocm():
pytest.skip(f"{kv_cache_dtype} is currently not supported on ROCm/HIP.")
if kv_cache_dtype == "fp8_e5m2" and current_platform.is_cuda():
pytest.skip(f"{kv_cache_dtype} is not supported by FLASH_ATTN on CUDA.")
if not (
current_platform.is_xpu()
tests/parser/engine/test_qwen3.py
+2 -2
View File
@@ -454,10 +454,10 @@ class TestStreaming:
args_after_partial_tag = collect_tool_arguments(results[:4])
assert "<param" not in args_after_partial_tag
assert args_after_partial_tag == '{"query": "hello'
assert args_after_partial_tag == '{"query": "hello '
args_text = collect_tool_arguments(results)
assert json.loads(args_text) == {"query": "hello", "limit": "10"}
assert json.loads(args_text) == {"query": "hello ", "limit": "10"}
def test_streaming_numeric_values(self, parser, mock_request):
chunks = [
tests/quantization/test_fp8.py
+3
View File
@@ -93,6 +93,9 @@ def test_online_quantization(
use_rocm_aiter: bool,
monkeypatch,
) -> None:
if kv_cache_dtype == "fp8" and current_platform.is_device_capability_family(90):
pytest.skip("FA3 currently rejects FP8 KV cache output dtype on SM90")
if use_rocm_aiter:
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
tests/quantization/test_per_token_kv_cache.py
+4 -1
View File
@@ -717,7 +717,10 @@ def test_triton_unified_attention_per_token_head_scale(
# Coarser quantization → wider tolerance.
if is_int4:
atol, rtol = 0.5, 0.5
# Hopper's attention reduction order can move a few BF16 elements by
# just over 1.0 after INT4 quantization.
atol = 1.1 if current_platform.is_device_capability_family(90) else 0.5
rtol = 0.5
else:
atol, rtol = 5e-2, 5e-2
torch.testing.assert_close(output_q, output_ref, atol=atol, rtol=rtol)
tests/test_pooling_params.py
+38
View File
@@ -1,12 +1,18 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
from typing import Any
import pytest
from pydantic import TypeAdapter, ValidationError
from tests.models.utils import EmbedModelInfo
from vllm import PoolingParams
from vllm.config import ModelConfig, PoolerConfig
from vllm.entrypoints.pooling.classify.protocol import ClassificationRequest
from vllm.entrypoints.pooling.embed.protocol import EmbeddingRequest
from vllm.entrypoints.pooling.pooling.protocol import PoolingRequest
from vllm.exceptions import VLLMValidationError
EMBEDDING_MODELS = [
EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False),
@@ -27,6 +33,38 @@ class MockModelConfig:
pooler_config: PoolerConfig
@pytest.mark.parametrize(
("parameter", "value", "message"),
[
(
"normalize",
False,
"Parameter `normalize` was removed; use `use_activation` instead.",
),
("task", "score", "`score` task was removed; use `classify` instead."),
(
"task",
"encode",
"`encode` task was removed; use `token_embed` or `token_classify` instead.",
),
],
)
def test_removed_pooling_parameters(parameter: str, value: Any, message: str):
data = {"input": "hello", parameter: value}
for request_type in (EmbeddingRequest, ClassificationRequest, PoolingRequest):
with pytest.raises(ValidationError, match=message) as exc_info:
TypeAdapter(request_type).validate_python(data)
assert len(exc_info.value.errors()) == 1
with pytest.raises(ValidationError, match=message) as exc_info:
TypeAdapter(PoolerConfig).validate_python({parameter: value})
assert len(exc_info.value.errors()) == 1
if parameter == "task":
with pytest.raises(VLLMValidationError, match=message):
PoolingParams(task=value)
def test_embed():
task = "embed"
model_config = MockModelConfig(pooler_config=PoolerConfig(seq_pooling_type="CLS"))
tests/tokenizers_/test_detokenize.py
+120
View File
@@ -8,6 +8,7 @@ import pytest
from transformers import AutoTokenizer, PythonBackend, TokenizersBackend
from vllm.sampling_params import SamplingParams
from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens
from vllm.tokenizers.mistral import MistralTokenizer
from vllm.v1.engine import EngineCoreRequest
from vllm.v1.engine.detokenizer import (
@@ -239,3 +240,122 @@ def test_oov_decode(tokenizer, fast):
assert decoded_text == ""
assert out_ids == [len(tokenizer)]
# ---------- convert_ids_list_to_tokens collision tests ----------
class _MockBackend:
"""Fake backend_tokenizer that exposes pre_tokenizer config."""
def __init__(self, pre_tokenizer_type, replacement=None):
import json
pre: dict = {"type": pre_tokenizer_type}
if replacement is not None:
pre["replacement"] = replacement
self._config = json.dumps({"pre_tokenizer": pre})
def to_str(self):
return self._config
class _MockTokenizer:
"""Minimal tokenizer mock for testing convert_ids_list_to_tokens."""
def __init__(
self,
raw_tokens: dict[int, str],
decoded_tokens: dict[int, str],
pre_tokenizer_type: str = "Metaspace",
replacement: str | None = "▁",
):
self._raw = raw_tokens
self._decoded = decoded_tokens
self.backend_tokenizer = _MockBackend(pre_tokenizer_type, replacement)
def convert_ids_to_tokens(
self, ids: list[int], skip_special_tokens: bool = False
) -> list[str]:
return [self._raw[tid] for tid in ids]
def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
return "".join(self._decoded[tid] for tid in ids)
def test_sentencepiece_leading_space_preserved():
"""▁true and true must produce distinct strings."""
tok = _MockTokenizer(
raw_tokens={0: "▁true", 1: "true", 2: "▁false", 3: "false"},
decoded_tokens={0: "true", 1: "true", 2: "false", 3: "false"},
)
result = convert_ids_list_to_tokens(tok, [0, 1, 2, 3])
assert result == [" true", "true", " false", "false"]
# No dict collision when used as top_logprobs keys
logprobs = dict(zip(result, [-0.1, -0.2, -0.3, -0.4]))
assert len(logprobs) == 4
def test_whitespace_run_tokens_stay_distinct():
"""▁, ▁▁, ▁▁▁ must produce different-length space strings."""
tok = _MockTokenizer(
raw_tokens={0: "▁", 1: "▁▁", 2: "▁▁▁"},
decoded_tokens={0: "", 1: " ", 2: " "},
)
result = convert_ids_list_to_tokens(tok, [0, 1, 2])
assert result == [" ", " ", " "]
def test_bpe_leading_space_already_preserved():
"""GPT-2 BPE: Ġtrue already decodes to ' true', no fix needed."""
tok = _MockTokenizer(
raw_tokens={0: "Ġtrue", 1: "true"},
decoded_tokens={0: " true", 1: "true"},
pre_tokenizer_type="ByteLevel",
replacement=None,
)
result = convert_ids_list_to_tokens(tok, [0, 1])
assert result == [" true", "true"]
def test_logprobs_count_stable_across_k():
"""logprobs=4 and logprobs=10 must return 4 and 10 entries."""
tok = _MockTokenizer(
raw_tokens={
0: "▁true",
1: "a",
2: "b",
3: "c",
4: "true",
5: "d",
6: "e",
7: "f",
8: "g",
9: "h",
},
decoded_tokens={
0: "true",
1: "a",
2: "b",
3: "c",
4: "true",
5: "d",
6: "e",
7: "f",
8: "g",
9: "h",
},
)
ids = list(range(10))
lps = [-0.1 * (i + 1) for i in range(10)]
tokens4 = convert_ids_list_to_tokens(tok, ids[:4])
top4 = dict(zip(tokens4, lps[:4]))
tokens10 = convert_ids_list_to_tokens(tok, ids)
top10 = dict(zip(tokens10, lps))
assert len(top4) == 4
assert len(top10) == 10
assert top4[" true"] == top10[" true"]
tests/tool_parsers/test_minicpm5xml_tool_parser.py
+52
View File
@@ -733,3 +733,55 @@ def test_extract_tool_calls_streaming_multiple(parser: ToolParser) -> None:
"nums": [7, 8, 9],
"exact": False,
}
def make_tools_write() -> list[ChatCompletionToolsParam]:
return [
_tool(
"write_file",
{
"type": "object",
"properties": {"content": {"type": "string"}},
"required": ["content"],
},
)
]
class TestParameterWhitespace:
"""CDATA is verbatim, so its whitespace must survive."""
def test_cdata_whitespace_preserved(self, parser: ToolParser) -> None:
request = make_request(make_tools_write())
text = (
'<function name="write_file">'
'<param name="content"><![CDATA[ def foo():\n pass\n]]></param>'
"</function>\n"
)
out = parser.extract_tool_calls(text, request)
assert len(out.tool_calls) == 1
args = json.loads(out.tool_calls[0].function.arguments)
assert args["content"] == " def foo():\n pass\n"
def test_cdata_whitespace_preserved_streaming(self, parser: ToolParser) -> None:
"""Value split across chunks, i.e. the partial path."""
request = make_request(make_tools_write())
chunks = [
'<function name="write_file">',
'<param name="content"><![CDATA[ def foo():\n',
" pass\n]]></param></function>\n",
]
reconstructor = run_tool_extraction_streaming(
parser,
chunks,
request,
assert_one_tool_per_delta=False,
)
assert len(reconstructor.tool_calls) == 1
assert json.loads(reconstructor.tool_calls[0].function.arguments) == {
"content": " def foo():\n pass\n"
}
tests/tool_parsers/test_minimax_m2_tool_parser.py
+34
View File
@@ -567,3 +567,37 @@ class TestNoneStringPreservation:
assert len(tc) == 1
parsed = json.loads(tc[0]["arguments"])
assert parsed["value"] == "nil"
class TestParameterWhitespace:
"""Parameter values must preserve surrounding whitespace."""
def test_whitespace_preserved(self, parser):
results = _feed(
parser,
[
'<minimax:tool_call><invoke name="echo">'
'<parameter name="msg"> hi </parameter>'
"</invoke></minimax:tool_call>",
],
)
tc = _collect_tool_calls(results)
assert len(tc) == 1
assert json.loads(tc[0]["arguments"]) == {"msg": " hi "}
def test_whitespace_preserved_across_chunks(self, parser):
"""Value split before </parameter> arrives, i.e. the partial path."""
results = _feed(
parser,
[
'<minimax:tool_call><invoke name="write">'
'<parameter name="content"> def foo():\n',
" return 1\n",
"</parameter></invoke></minimax:tool_call>",
],
)
tc = _collect_tool_calls(results)
assert len(tc) == 1
assert json.loads(tc[0]["arguments"]) == {
"content": " def foo():\n return 1\n"
}
tests/tool_parsers/test_qwen3coder_tool_parser.py
+55
View File
@@ -1482,3 +1482,58 @@ def test_adjust_request_required_prefers_structural_tag(
out = TestParser(MagicMock(), tools=sample_tools).adjust_request(req)
assert out.structured_outputs is not None
assert out.structured_outputs.structural_tag is not None
WRITE_FILE_TOOLS = [
ChatCompletionToolsParam(
type="function",
function={
"name": "write_file",
"parameters": {
"type": "object",
"properties": {"content": {"type": "string"}},
},
},
)
]
WRITE_FILE_OUTPUT = (
"<tool_call>\n<function=write_file>\n"
"<parameter=content>\n"
" def foo():\n return 1\n\n"
"</parameter>\n</function>\n</tool_call>"
)
EXPECTED_CONTENT = " def foo():\n return 1\n"
class TestParameterWhitespace:
"""Wrapping newlines are markup; the rest of the value is preserved."""
def test_whitespace_preserved(self, qwen3_tokenizer):
parser = Qwen3EngineToolParser(qwen3_tokenizer, tools=WRITE_FILE_TOOLS)
request = ChatCompletionRequest(
model=MODEL, messages=[], tools=WRITE_FILE_TOOLS
)
extracted = parser.extract_tool_calls(WRITE_FILE_OUTPUT, request=request)
args = json.loads(extracted.tool_calls[0].function.arguments)
assert args["content"] == EXPECTED_CONTENT
def test_whitespace_preserved_streaming(self, qwen3_tokenizer):
"""The partial path must not strip the value either."""
parser = Qwen3EngineToolParser(qwen3_tokenizer, tools=WRITE_FILE_TOOLS)
request = ChatCompletionRequest(
model=MODEL, messages=[], tools=WRITE_FILE_TOOLS
)
streamed = ""
for delta_message in stream_delta_message_generator(
parser, qwen3_tokenizer, WRITE_FILE_OUTPUT, request
):
for tool_call in delta_message.tool_calls or []:
if tool_call.function and tool_call.function.arguments:
streamed += tool_call.function.arguments
assert json.loads(streamed)["content"] == EXPECTED_CONTENT
tests/utils.py
+14 -1
View File
@@ -108,6 +108,7 @@ if current_platform.is_rocm():
elif current_platform.is_cuda():
from vllm.third_party.pynvml import (
nvmlDeviceGetHandleByIndex,
nvmlDeviceGetHandleByUUID,
nvmlDeviceGetMemoryInfo,
nvmlInit,
nvmlShutdown,
@@ -1521,6 +1522,18 @@ def get_physical_device_indices(devices: list[int]):
return [index_mapping[i] for i in devices if i in index_mapping]
def get_nvml_device_handle(device: int):
visible_devices = os.environ.get("NVIDIA_VISIBLE_DEVICES")
if visible_devices is not None:
identifiers = visible_devices.split(",")
if device < len(identifiers):
identifier = identifiers[device]
if identifier.startswith(("GPU-", "MIG-")):
return nvmlDeviceGetHandleByUUID(identifier)
return nvmlDeviceGetHandleByIndex(device)
@_nvml()
def record_gpu_memory_usage_stats(
*,
@@ -1534,7 +1547,7 @@ def record_gpu_memory_usage_stats(
gb_used = mem_info["vram_used"] / 2**10
gb_total = mem_info["vram_total"] / 2**10
else:
dev_handle = nvmlDeviceGetHandleByIndex(device)
dev_handle = get_nvml_device_handle(device)
mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
gb_used = mem_info.used / 2**30
gb_total = mem_info.total / 2**30
tests/v1/attention/test_attention_splitting.py
+16
View File
@@ -160,6 +160,8 @@ def apply_split_decodes_and_prefills(
decode_threshold: int,
require_uniform: bool,
padded_num_tokens: int | None = None,
is_prefilling: list[bool] | None = None,
treat_short_extends_as_decodes: bool = True,
):
"""Helper function to apply split_decodes_and_prefills and return
the results."""
@@ -173,11 +175,14 @@ def apply_split_decodes_and_prefills(
if padded_num_tokens is not None:
common_metadata.num_actual_tokens = padded_num_tokens
if is_prefilling is not None:
common_metadata.is_prefilling = torch.tensor(is_prefilling)
return split_decodes_and_prefills(
common_metadata,
decode_threshold=decode_threshold,
require_uniform=require_uniform,
treat_short_extends_as_decodes=treat_short_extends_as_decodes,
)
@@ -236,6 +241,17 @@ def test_split_decodes_and_prefills_uniform_all_ones():
assert num_prefill_tokens == 0
def test_split_decodes_and_prefills_uniform_short_extend():
result = apply_split_decodes_and_prefills(
[1, 1],
decode_threshold=1,
require_uniform=True,
is_prefilling=[False, True],
treat_short_extends_as_decodes=False,
)
assert result == (1, 1, 1, 1)
def test_split_decodes_and_prefills_uniform_all_short_decodes():
query_lens = [2, 2, 1, 3, 2, 1, 2]
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
tests/v1/attention/test_backend_per_kind.py
+67
View File
@@ -0,0 +1,67 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for per-KV-group attention backend selection (backend_per_kind)."""
import pytest
from vllm.config.attention import AttentionConfig
from vllm.v1.attention.backend import AttentionType
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.attention.selector import get_attn_spec_kind
from vllm.v1.kv_cache_interface import KVCacheSpecKind
@pytest.mark.parametrize(
"signals,expected",
[
(dict(use_mla=False, has_sliding_window=False), "full"),
(dict(use_mla=True, has_sliding_window=False), "mla"),
(dict(use_mla=True, has_sliding_window=True), "sw_mla"),
(dict(use_mla=False, has_sliding_window=True), "sw"),
],
)
def test_get_attn_spec_kind_decoder(signals, expected):
kind_by_name = {
"full": KVCacheSpecKind.FULL_ATTENTION,
"mla": KVCacheSpecKind.MLA_ATTENTION,
"sw_mla": KVCacheSpecKind.SLIDING_WINDOW_MLA,
"sw": KVCacheSpecKind.SLIDING_WINDOW,
}
kind = get_attn_spec_kind(attn_type=AttentionType.DECODER, **signals)
assert kind is kind_by_name[expected]
@pytest.mark.parametrize(
"attn_type,expected",
[
(AttentionType.ENCODER_ONLY, KVCacheSpecKind.ENCODER_ONLY_ATTENTION),
(AttentionType.ENCODER_DECODER, KVCacheSpecKind.CROSS_ATTENTION),
],
)
def test_get_attn_spec_kind_attn_type(attn_type, expected):
kind = get_attn_spec_kind(
use_mla=False,
has_sliding_window=False,
attn_type=attn_type,
)
assert kind is expected
def test_backend_per_kind_parses_strings():
cfg = AttentionConfig(
backend_per_kind={
"mla_attention": "FLASHINFER_MLA",
"sliding_window_mla": "triton_mla", # case-insensitive
}
)
assert cfg.backend_per_kind["mla_attention"] is AttentionBackendEnum.FLASHINFER_MLA
assert cfg.backend_per_kind["sliding_window_mla"] is AttentionBackendEnum.TRITON_MLA
def test_backend_per_kind_rejects_unknown_kind():
with pytest.raises(ValueError, match="Unknown KV cache group kind"):
AttentionConfig(backend_per_kind={"not_a_kind": "TRITON_MLA"})
def test_backend_per_kind_defaults_empty():
assert AttentionConfig().backend_per_kind == {}
tests/v1/attention/test_sparse_mla_backends.py
+7
View File
@@ -1114,6 +1114,13 @@ def test_sparse_backend_prefill_correctness(
@pytest.mark.parametrize(
"seq_lens,query_lens,workspace_size,max_logits_bytes,expected",
[
(
torch.tensor([0]),
torch.tensor([0]),
100,
1000,
[],
),
# Logits constraint triggers split (M*N exceeds budget)
# req0: M=10, N=100 -> 1000 elems (4000 bytes) - fits in 5000
# req1: adding M=10, N=100 -> new_M=20, new_N=200 -> 4000 elems > 1250
tests/v1/core/test_async_scheduler.py
+1
View File
@@ -280,6 +280,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
scheduler.waiting = Mock()
scheduler.kv_cache_manager = Mock()
scheduler.kv_cache_manager.take_events.return_value = None
scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
scheduler.kv_event_publisher = Mock()
scheduler.finished_req_ids = set()
scheduler.finished_req_ids_dict = None
tests/v1/core/test_scheduler.py
+1
View File
@@ -3036,6 +3036,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
scheduler.waiting = Mock()
scheduler.kv_cache_manager = Mock()
scheduler.kv_cache_manager.take_events.return_value = None
scheduler.kv_cache_manager.estimate_cached_tokens.return_value = 0
scheduler.kv_event_publisher = Mock()
scheduler.finished_req_ids = set()
scheduler.finished_req_ids_dict = None
tests/v1/e2e/general/test_attention_backend_per_kind.py
+94
View File
@@ -0,0 +1,94 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end checks that ``backend_per_kind`` selects the requested attention
backend for each KV-cache group at runtime.
Uses ``google/gemma-3-1b-it``, which interleaves full-attention and
sliding-window layers, so the model produces separate ``full_attention`` and
``sliding_window`` KV-cache groups.
"""
import pytest
from vllm import LLM
from vllm.config.attention import AttentionConfig
from vllm.platforms import current_platform
MODEL = "google/gemma-3-1b-it"
def _collect_group_backends(worker) -> list[tuple[str, str]]:
"""Runs on the worker: returns (spec_kind, backend_name) per attn group."""
from vllm.v1.kv_cache_interface import get_kv_cache_spec_kind
out: list[tuple[str, str]] = []
for kv_group in worker.model_runner.attn_groups:
for attn_group in kv_group:
kind = get_kv_cache_spec_kind(attn_group.kv_cache_spec)
out.append((kind.value, attn_group.backend.get_name()))
return out
@pytest.mark.skipif(
not current_platform.is_cuda(), reason="backend names are CUDA-specific"
)
@pytest.mark.parametrize(
"backend_per_kind",
[
{"full_attention": "FLASH_ATTN", "sliding_window": "TRITON_ATTN"},
# Swapped, to prove the mapping is causal rather than the default.
{"full_attention": "TRITON_ATTN", "sliding_window": "FLASH_ATTN"},
],
)
def test_backend_per_kind_splits_groups(backend_per_kind, monkeypatch):
# collective_rpc ships the callable to the EngineCore subprocess; the
# secure msgpack encoder can't serialize functions, so opt into the
# pickle fallback (same pattern as test_pooling_chunked_prefill).
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
llm = LLM(
model=MODEL,
attention_config=AttentionConfig(backend_per_kind=backend_per_kind),
enforce_eager=True,
max_model_len=2048,
gpu_memory_utilization=0.4,
)
group_backends = llm.llm_engine.collective_rpc(_collect_group_backends)[0]
kinds = {kind for kind, _ in group_backends}
# gemma3 must actually split into both kinds for this test to be meaningful.
assert "full_attention" in kinds
assert "sliding_window" in kinds
for kind, backend_name in group_backends:
if kind in backend_per_kind:
assert backend_name == backend_per_kind[kind], (
f"{kind} group used {backend_name}, expected {backend_per_kind[kind]}"
)
@pytest.mark.skipif(
not current_platform.is_cuda(), reason="backend names are CUDA-specific"
)
def test_backend_per_kind_overrides_global_backend(monkeypatch):
"""A per-kind entry wins over the global ``backend`` for its kind; other
kinds fall back to the global backend."""
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
llm = LLM(
model=MODEL,
attention_config=AttentionConfig(
backend="FLASH_ATTN",
backend_per_kind={"sliding_window": "TRITON_ATTN"},
),
enforce_eager=True,
max_model_len=2048,
gpu_memory_utilization=0.4,
)
group_backends = llm.llm_engine.collective_rpc(_collect_group_backends)[0]
for kind, backend_name in group_backends:
if kind == "sliding_window":
assert backend_name == "TRITON_ATTN"
elif kind == "full_attention":
assert backend_name == "FLASH_ATTN"
tests/v1/e2e/general/test_correctness_sliding_window.py
+28 -23
View File
@@ -33,7 +33,7 @@ model_config = {
@pytest.mark.parametrize("seed", [1])
@pytest.mark.parametrize("disable_hybrid_kv_cache_manager", [True, False])
def test_sliding_window_retrieval(
model, batch_size, seed, disable_hybrid_kv_cache_manager
model, batch_size, seed, disable_hybrid_kv_cache_manager, vllm_runner
):
"""
The test does a bunch of assignments "x1 = 10\nx2 = 33\n..." and then
@@ -48,34 +48,39 @@ def test_sliding_window_retrieval(
test_config = model_config[model]
llm = LLM(
model=model,
with vllm_runner(
model,
max_model_len=None,
enable_chunked_prefill=None,
disable_hybrid_kv_cache_manager=disable_hybrid_kv_cache_manager,
enforce_eager=enforce_eager,
)
sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
) as runner:
llm = runner.get_llm()
sampling_params = SamplingParams(temperature=0.0, max_tokens=100)
prompts, answer, indices = prep_prompts(batch_size, ln_range=test_config.ln_range)
prompts, answer, indices = prep_prompts(
batch_size, ln_range=test_config.ln_range
)
check_length(prompts, llm, test_config.sliding_window)
check_length(prompts, llm, test_config.sliding_window)
# Fresh generation
responses = llm.generate(prompts, sampling_params)
check_answers(
indices,
answer,
[response.outputs[0].text for response in responses],
accept_rate=1.0,
)
# Fresh generation
responses = llm.generate(prompts, sampling_params)
check_answers(
indices,
answer,
[response.outputs[0].text for response in responses],
accept_rate=1.0,
)
# Re-generate with the same prompts to test prefix caching
responses = llm.generate(prompts, sampling_params)
check_answers(
indices,
answer,
[response.outputs[0].text for response in responses],
accept_rate=1.0,
)
# Re-generate with the same prompts to test prefix caching
responses = llm.generate(prompts, sampling_params)
check_answers(
indices,
answer,
[response.outputs[0].text for response in responses],
accept_rate=1.0,
)
def check_length(prompts: list[str], llm: LLM, sliding_window: int):
tests/v1/kv_connector/unit/test_moriio_proxy_routing.py
+205
View File
@@ -0,0 +1,205 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Hardware-fair request routing in the MoRIIO toy P/D proxy.
Exercises the REAL ``flat_interleaved_dp_route`` from the toy proxy (loaded from
its ``examples/`` path with heavy deps stubbed) — no routing logic is copied
here. This is the request-distribution half of end-to-end fairness: the proxy
must hand every prefill/decode (instance, dp_rank) slot an equal share so no GPU
is starved. The connector-side read routing is covered in
``test_moriio_routing_fairness.py``.
Regression target: the previous scheme derived ``instance = req % n_instances``
and ``dp_rank = req % dp_size`` from the same counter, so when
``n_instances | dp_size`` each instance was locked to a stride-``n`` subset of
its ranks — e.g. 2 prefill instances x DP8 stranded 4 of every node's 8 GPUs.
``flat_interleaved_dp_route`` walks ONE counter over the full
``(instance, dp_rank)`` slot space, so the two selections can never alias.
Role shapes below are exactly those in the RFC (#46107) deployments:
(1, 1) 1P/1D TP8 (2, 1) 2P/2D TP8 (4, 1) 4D TP8
(1, 8) 1D DP8EP (2, 8) 2P DP8EP (3, 8) 3D DP8EP
"""
import contextlib
import importlib.util
import sys
import types
from collections import Counter
from pathlib import Path
from typing import cast
import pytest
_MISSING = object()
PROXY_REL = "examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py"
def _module(name, **attrs):
module = types.ModuleType(name)
for attr, value in attrs.items():
setattr(module, attr, value)
return module
def _package(name):
module = _module(name)
module.__path__ = []
return module
class _QuartStub:
def __init__(self, *a, **k):
pass
def route(self, *a, **k):
return lambda fn: fn
def post(self, *a, **k):
return self.route()
async def _make_response_stub(value):
return value
@contextlib.contextmanager
def _proxy_import_stubs():
"""Stub the proxy's external deps so the module imports for a pure unit test.
Only third-party/vllm imports are stubbed; the routing function under test
is executed as-is from the real module.
"""
common = "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common"
class _MoRIIOConstants:
TRANSFER_PREFIX = "moriio-transfer"
stubs = {
"aiohttp": _module("aiohttp"),
"msgpack": _module("msgpack"),
"zmq": _module("zmq"),
"quart": _module(
"quart",
Quart=_QuartStub,
Request=object,
make_response=_make_response_stub,
request=object(),
),
"vllm": _package("vllm"),
"vllm.distributed": _package("vllm.distributed"),
"vllm.distributed.kv_transfer": _package("vllm.distributed.kv_transfer"),
"vllm.distributed.kv_transfer.kv_connector": _package(
"vllm.distributed.kv_transfer.kv_connector"
),
"vllm.distributed.kv_transfer.kv_connector.v1": _package(
"vllm.distributed.kv_transfer.kv_connector.v1"
),
"vllm.distributed.kv_transfer.kv_connector.v1.moriio": _package(
"vllm.distributed.kv_transfer.kv_connector.v1.moriio"
),
common: _module(common, MoRIIOConstants=_MoRIIOConstants),
}
saved = {}
for name, module in stubs.items():
if name not in sys.modules:
saved[name] = _MISSING
sys.modules[name] = module
try:
yield
finally:
for name, previous in saved.items():
if previous is _MISSING:
sys.modules.pop(name, None)
else:
sys.modules[name] = cast(types.ModuleType, previous)
def _load_proxy_module():
path = Path(__file__).parents[4] / PROXY_REL
spec = importlib.util.spec_from_file_location("moriio_proxy_under_test", path)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
with _proxy_import_stubs():
spec.loader.exec_module(module)
return module
@pytest.fixture(scope="module")
def route():
return _load_proxy_module().flat_interleaved_dp_route
def _instances(n: int, dp_size: int):
# Only dp_size is read by the router; tp_size carried for realism.
return [{"dp_size": dp_size, "tp_size": 8 // dp_size} for _ in range(n)]
# (n_instances, dp_size) for every distinct P/D role shape in the RFC configs.
ROLE_SHAPES = [(1, 1), (2, 1), (4, 1), (1, 8), (2, 8), (3, 8)]
SHAPE_IDS = [f"n{n}_dp{dp}" for n, dp in ROLE_SHAPES]
def _route_n(route, instances, count):
# Proxy uses 1-indexed request numbers (slot = (request_number - 1) % ...).
return [route(rn, instances) for rn in range(1, count + 1)]
@pytest.mark.parametrize(("n", "dp"), ROLE_SHAPES, ids=SHAPE_IDS)
def test_full_slot_space_is_covered_uniformly(route, n, dp):
instances = _instances(n, dp)
period = n * dp
# Three full cycles -> every (instance, dp_rank) slot must be hit the same
# number of times (exactly uniform, no starved slot, no aliasing).
hits = Counter(_route_n(route, instances, period * 3))
expected: set[tuple[int, int | None]]
if dp == 1:
expected = {(inst, None) for inst in range(n)}
else:
expected = {(inst, r) for inst in range(n) for r in range(dp)}
assert set(hits) == expected, f"missing slots: {expected - set(hits)}"
assert max(hits.values()) == min(hits.values())
@pytest.mark.parametrize(("n", "dp"), ROLE_SHAPES, ids=SHAPE_IDS)
def test_instance_and_dp_rank_marginals_are_balanced(route, n, dp):
instances = _instances(n, dp)
routed = _route_n(route, instances, n * dp * 5)
inst_counts = Counter(inst for inst, _ in routed)
assert set(inst_counts) == set(range(n))
assert max(inst_counts.values()) == min(inst_counts.values())
dp_counts = Counter(r for _, r in routed)
if dp == 1:
assert set(dp_counts) == {None}
else:
assert set(dp_counts) == set(range(dp))
assert max(dp_counts.values()) == min(dp_counts.values())
def test_tp_instance_forwards_no_dp_rank(route):
# dp_size == 1 (a TP instance) must yield dp_rank None so the proxy never
# forwards an out-of-range data-parallel rank.
assert all(r is None for _, r in _route_n(route, _instances(2, 1), 8))
def test_two_instance_dp8_gives_every_node_all_ranks(route):
# Direct regression for the stranded-GPU bug: with 2 prefill instances x DP8
# each instance must receive ALL 8 dp ranks (16 distinct slots), not a
# stride-2 subset of 4.
routed = _route_n(route, _instances(2, 8), 2 * 8)
per_instance: dict[int, set] = {0: set(), 1: set()}
for inst, dp_rank in routed:
per_instance[inst].add(dp_rank)
assert per_instance[0] == set(range(8))
assert per_instance[1] == set(range(8))
def test_consecutive_requests_alternate_instances(route):
# Interleaved order spreads consecutive requests across instances rather
# than filling one instance's ranks before moving on.
insts = [inst for inst, _ in _route_n(route, _instances(3, 8), 3)]
assert insts == [0, 1, 2]
tests/v1/kv_offload/test_factory.py
+4 -4
View File
@@ -385,7 +385,7 @@ def test_tiering_spec_aligns_row_size():
assert spec.num_blocks == cpu_bytes_to_use // alignment
def test_offloading_spec_resolves_prefill_context_parallel_block_sizes():
def test_offloading_spec_kv_sharding_ignores_prefill_context_parallel():
config = _make_layout_vllm_config(
cpu_bytes_to_use=65536,
extra_config={"block_size": 64},
@@ -394,9 +394,9 @@ def test_offloading_spec_resolves_prefill_context_parallel_block_sizes():
spec = _create_spec(config, _make_kv_cache_config())
assert spec.tokens_per_block == (32,)
assert spec.tokens_per_hash == 32
assert spec.blocks_per_chunk == 2
assert spec.tokens_per_block == (16,)
assert spec.tokens_per_hash == 16
assert spec.blocks_per_chunk == 4
def test_offloading_config_preserves_data_parallel_index():

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