forked from Karylab-cklius/vllm
Compare commits
72
Commits
@@ -1,7 +1,8 @@
|
||||
name: vllm_ci
|
||||
job_dirs:
|
||||
- ".buildkite/test_areas"
|
||||
- ".buildkite/image_build"
|
||||
- ".buildkite/test_areas"
|
||||
- ".buildkite/hardware_tests"
|
||||
run_all_patterns:
|
||||
- "docker/Dockerfile"
|
||||
- "CMakeLists.txt"
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
group: Hardware
|
||||
steps:
|
||||
- label: "AMD: :docker: build image"
|
||||
device: amd_cpu
|
||||
no_plugin: true
|
||||
commands:
|
||||
- >
|
||||
docker build
|
||||
--build-arg max_jobs=16
|
||||
--build-arg REMOTE_VLLM=1
|
||||
--build-arg ARG_PYTORCH_ROCM_ARCH='gfx90a;gfx942'
|
||||
--build-arg VLLM_BRANCH=$BUILDKITE_COMMIT
|
||||
--tag "rocm/vllm-ci:${BUILDKITE_COMMIT}"
|
||||
-f docker/Dockerfile.rocm
|
||||
--target test
|
||||
--no-cache
|
||||
--progress plain .
|
||||
- docker push "rocm/vllm-ci:${BUILDKITE_COMMIT}"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: -1 # Agent was lost
|
||||
limit: 1
|
||||
- exit_status: -10 # Agent was lost
|
||||
limit: 1
|
||||
- exit_status: 1 # Machine occasionally fail
|
||||
limit: 1
|
||||
@@ -0,0 +1,8 @@
|
||||
group: Hardware
|
||||
steps:
|
||||
- label: "Arm CPU Test"
|
||||
soft_fail: true
|
||||
device: arm_cpu
|
||||
no_plugin: true
|
||||
commands:
|
||||
- bash .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
|
||||
@@ -0,0 +1,10 @@
|
||||
group: Hardware
|
||||
depends_on: ~
|
||||
steps:
|
||||
- label: "Ascend NPU Test"
|
||||
soft_fail: true
|
||||
timeout_in_minutes: 20
|
||||
no_plugin: true
|
||||
device: ascend_npu
|
||||
commands:
|
||||
- bash .buildkite/scripts/hardware_ci/run-npu-test.sh
|
||||
@@ -0,0 +1,10 @@
|
||||
group: Hardware
|
||||
steps:
|
||||
- label: "GH200 Test"
|
||||
soft_fail: true
|
||||
device: gh200
|
||||
no_plugin: true
|
||||
optional: true
|
||||
commands:
|
||||
- nvidia-smi
|
||||
- bash .buildkite/scripts/hardware_ci/run-gh200-test.sh
|
||||
@@ -0,0 +1,23 @@
|
||||
group: Hardware
|
||||
depends_on: ~
|
||||
steps:
|
||||
- label: "Intel CPU Test"
|
||||
soft_fail: true
|
||||
device: intel_cpu
|
||||
no_plugin: true
|
||||
commands:
|
||||
- bash .buildkite/scripts/hardware_ci/run-cpu-test.sh
|
||||
|
||||
- label: "Intel HPU Test"
|
||||
soft_fail: true
|
||||
device: intel_hpu
|
||||
no_plugin: true
|
||||
commands:
|
||||
- bash .buildkite/scripts/hardware_ci/run-hpu-test.sh
|
||||
|
||||
- label: "Intel GPU Test"
|
||||
soft_fail: true
|
||||
device: intel_gpu
|
||||
no_plugin: true
|
||||
commands:
|
||||
- bash .buildkite/scripts/hardware_ci/run-xpu-test.sh
|
||||
@@ -1,56 +1,254 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
set -euo pipefail
|
||||
|
||||
if [[ $# -lt 8 ]]; then
|
||||
echo "Usage: $0 <registry> <repo> <commit> <branch> <vllm_use_precompiled> <vllm_merge_base_commit> <cache_from> <cache_to>"
|
||||
exit 1
|
||||
# replace invalid characters in Docker image tags and truncate to 128 chars
|
||||
clean_docker_tag() {
|
||||
local input="$1"
|
||||
echo "$input" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128
|
||||
}
|
||||
|
||||
print_usage_and_exit() {
|
||||
echo "Usage: $0 <registry> <repo> <commit> <branch> <vllm_use_precompiled> <vllm_merge_base_commit> <cache_from> <cache_to>"
|
||||
exit 1
|
||||
}
|
||||
|
||||
print_instance_info() {
|
||||
echo ""
|
||||
echo "=== Debug: Instance Information ==="
|
||||
# Get IMDSv2 token
|
||||
if TOKEN=$(curl -s -X PUT "http://169.254.169.254/latest/api/token" \
|
||||
-H "X-aws-ec2-metadata-token-ttl-seconds: 21600" 2>/dev/null); then
|
||||
AMI_ID=$(curl -s -H "X-aws-ec2-metadata-token: $TOKEN" \
|
||||
http://169.254.169.254/latest/meta-data/ami-id 2>/dev/null || echo "unknown")
|
||||
INSTANCE_TYPE=$(curl -s -H "X-aws-ec2-metadata-token: $TOKEN" \
|
||||
http://169.254.169.254/latest/meta-data/instance-type 2>/dev/null || echo "unknown")
|
||||
INSTANCE_ID=$(curl -s -H "X-aws-ec2-metadata-token: $TOKEN" \
|
||||
http://169.254.169.254/latest/meta-data/instance-id 2>/dev/null || echo "unknown")
|
||||
AZ=$(curl -s -H "X-aws-ec2-metadata-token: $TOKEN" \
|
||||
http://169.254.169.254/latest/meta-data/placement/availability-zone 2>/dev/null || echo "unknown")
|
||||
echo "AMI ID: ${AMI_ID}"
|
||||
echo "Instance Type: ${INSTANCE_TYPE}"
|
||||
echo "Instance ID: ${INSTANCE_ID}"
|
||||
echo "AZ: ${AZ}"
|
||||
else
|
||||
echo "Not running on EC2 or IMDS not available"
|
||||
fi
|
||||
# Check for warm cache AMI (marker file baked into custom AMI)
|
||||
if [[ -f /etc/vllm-ami-info ]]; then
|
||||
echo "Cache: warm (custom vLLM AMI)"
|
||||
cat /etc/vllm-ami-info
|
||||
else
|
||||
echo "Cache: cold (standard AMI)"
|
||||
fi
|
||||
echo "==================================="
|
||||
echo ""
|
||||
}
|
||||
|
||||
setup_buildx_builder() {
|
||||
echo "--- :buildkite: Setting up buildx builder"
|
||||
if [[ -S "${BUILDKIT_SOCKET}" ]]; then
|
||||
# Custom AMI with standalone buildkitd - use remote driver for warm cache
|
||||
echo "✅ Found local buildkitd socket at ${BUILDKIT_SOCKET}"
|
||||
echo "Using remote driver to connect to buildkitd (warm cache available)"
|
||||
if docker buildx inspect baked-vllm-builder >/dev/null 2>&1; then
|
||||
echo "Using existing baked-vllm-builder"
|
||||
docker buildx use baked-vllm-builder
|
||||
else
|
||||
echo "Creating baked-vllm-builder with remote driver"
|
||||
docker buildx create \
|
||||
--name baked-vllm-builder \
|
||||
--driver remote \
|
||||
--use \
|
||||
"unix://${BUILDKIT_SOCKET}"
|
||||
fi
|
||||
docker buildx inspect --bootstrap
|
||||
elif docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
|
||||
# Existing builder available
|
||||
echo "Using existing builder: ${BUILDER_NAME}"
|
||||
docker buildx use "${BUILDER_NAME}"
|
||||
docker buildx inspect --bootstrap
|
||||
else
|
||||
# No local buildkitd, no existing builder - create new docker-container builder
|
||||
echo "No local buildkitd found, using docker-container driver"
|
||||
docker buildx create --name "${BUILDER_NAME}" --driver docker-container --use
|
||||
docker buildx inspect --bootstrap
|
||||
fi
|
||||
|
||||
# builder info
|
||||
echo "Active builder:"
|
||||
docker buildx ls | grep -E '^\*|^NAME' || docker buildx ls
|
||||
}
|
||||
|
||||
check_and_skip_if_image_exists() {
|
||||
if [[ -n "${IMAGE_TAG:-}" ]]; then
|
||||
echo "--- :mag: Checking if image exists"
|
||||
if docker manifest inspect "${IMAGE_TAG}" >/dev/null 2>&1; then
|
||||
echo "Image already exists: ${IMAGE_TAG}"
|
||||
echo "Skipping build"
|
||||
exit 0
|
||||
fi
|
||||
echo "Image not found, proceeding with build"
|
||||
fi
|
||||
}
|
||||
|
||||
ecr_login() {
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
|
||||
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
|
||||
}
|
||||
|
||||
prepare_cache_tags() {
|
||||
# resolve and set: CACHE_TO, CACHE_FROM, CACHE_FROM_BASE_BRANCH, CACHE_FROM_MAIN
|
||||
TEST_CACHE_ECR="936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache"
|
||||
MAIN_CACHE_ECR="936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache"
|
||||
|
||||
if [[ "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
|
||||
if [[ "$BUILDKITE_BRANCH" == "main" ]]; then
|
||||
cache="${MAIN_CACHE_ECR}:latest"
|
||||
else
|
||||
clean_branch=$(clean_docker_tag "$BUILDKITE_BRANCH")
|
||||
cache="${TEST_CACHE_ECR}:${clean_branch}"
|
||||
fi
|
||||
CACHE_TO="$cache"
|
||||
CACHE_FROM="$cache"
|
||||
CACHE_FROM_BASE_BRANCH="$cache"
|
||||
else
|
||||
CACHE_TO="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}"
|
||||
CACHE_FROM="${TEST_CACHE_ECR}:pr-${BUILDKITE_PULL_REQUEST}"
|
||||
if [[ "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" == "main" ]]; then
|
||||
CACHE_FROM_BASE_BRANCH="${MAIN_CACHE_ECR}:latest"
|
||||
else
|
||||
clean_base=$(clean_docker_tag "$BUILDKITE_PULL_REQUEST_BASE_BRANCH")
|
||||
CACHE_FROM_BASE_BRANCH="${TEST_CACHE_ECR}:${clean_base}"
|
||||
fi
|
||||
fi
|
||||
|
||||
CACHE_FROM_MAIN="${MAIN_CACHE_ECR}:latest"
|
||||
export CACHE_TO CACHE_FROM CACHE_FROM_BASE_BRANCH CACHE_FROM_MAIN
|
||||
}
|
||||
|
||||
resolve_parent_commit() {
|
||||
if [[ -z "${PARENT_COMMIT:-}" ]]; then
|
||||
PARENT_COMMIT=$(git rev-parse HEAD~1 2>/dev/null || echo "")
|
||||
if [[ -n "${PARENT_COMMIT}" ]]; then
|
||||
echo "Computed parent commit for cache fallback: ${PARENT_COMMIT}"
|
||||
export PARENT_COMMIT
|
||||
else
|
||||
echo "Could not determine parent commit (may be first commit in repo)"
|
||||
fi
|
||||
else
|
||||
echo "Using provided PARENT_COMMIT: ${PARENT_COMMIT}"
|
||||
fi
|
||||
}
|
||||
|
||||
print_bake_config() {
|
||||
echo "--- :page_facing_up: Resolved bake configuration"
|
||||
BAKE_CONFIG_FILE="bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
|
||||
docker buildx bake -f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}" --print "${TARGET}" | tee "${BAKE_CONFIG_FILE}" || true
|
||||
echo "Saved bake config to ${BAKE_CONFIG_FILE}"
|
||||
echo "--- :arrow_down: Uploading bake config to Buildkite"
|
||||
buildkite-agent artifact upload "${BAKE_CONFIG_FILE}"
|
||||
}
|
||||
|
||||
#################################
|
||||
# Main Script #
|
||||
#################################
|
||||
print_instance_info
|
||||
|
||||
if [[ $# -lt 7 ]]; then
|
||||
print_usage_and_exit
|
||||
fi
|
||||
|
||||
# input args
|
||||
REGISTRY=$1
|
||||
REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
BRANCH=$4
|
||||
VLLM_USE_PRECOMPILED=$5
|
||||
VLLM_MERGE_BASE_COMMIT=$6
|
||||
CACHE_FROM=$7
|
||||
CACHE_TO=$8
|
||||
IMAGE_TAG=$7
|
||||
IMAGE_TAG_LATEST=${8:-} # only used for main branch, optional
|
||||
|
||||
# authenticate with AWS ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
|
||||
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com
|
||||
# build config
|
||||
TARGET="test-ci"
|
||||
CI_HCL_URL="${CI_HCL_URL:-https://raw.githubusercontent.com/vllm-project/ci-infra/main/docker/ci.hcl}"
|
||||
VLLM_BAKE_FILE="${VLLM_BAKE_FILE:-docker/docker-bake.hcl}"
|
||||
BUILDER_NAME="${BUILDER_NAME:-vllm-builder}"
|
||||
CI_HCL_PATH="/tmp/ci.hcl"
|
||||
BUILDKIT_SOCKET="/run/buildkit/buildkitd.sock"
|
||||
|
||||
# docker buildx
|
||||
docker buildx create --name vllm-builder --driver docker-container --use
|
||||
docker buildx inspect --bootstrap
|
||||
docker buildx ls
|
||||
prepare_cache_tags
|
||||
ecr_login
|
||||
|
||||
# skip build if image already exists
|
||||
if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT) ]]; then
|
||||
echo "Image not found, proceeding with build..."
|
||||
else
|
||||
echo "Image found"
|
||||
exit 0
|
||||
# Environment info (for docs and human readers)
|
||||
# CI_HCL_URL - URL to ci.hcl (default: from ci-infra main branch)
|
||||
# VLLM_CI_BRANCH - ci-infra branch to use (default: main)
|
||||
# VLLM_BAKE_FILE - Path to vLLM's bake file (default: docker/docker-bake.hcl)
|
||||
# BUILDER_NAME - Name for buildx builder (default: vllm-builder)
|
||||
#
|
||||
# Build configuration (exported as environment variables for bake):
|
||||
export BUILDKITE_COMMIT
|
||||
export PARENT_COMMIT
|
||||
export IMAGE_TAG
|
||||
export IMAGE_TAG_LATEST
|
||||
export CACHE_FROM
|
||||
export CACHE_FROM_BASE_BRANCH
|
||||
export CACHE_FROM_MAIN
|
||||
export CACHE_TO
|
||||
export VLLM_USE_PRECOMPILED
|
||||
export VLLM_MERGE_BASE_COMMIT
|
||||
|
||||
# print args
|
||||
echo "--- :mag: Arguments"
|
||||
echo "REGISTRY: ${REGISTRY}"
|
||||
echo "REPO: ${REPO}"
|
||||
echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
|
||||
echo "BRANCH: ${BRANCH}"
|
||||
echo "VLLM_USE_PRECOMPILED: ${VLLM_USE_PRECOMPILED}"
|
||||
echo "VLLM_MERGE_BASE_COMMIT: ${VLLM_MERGE_BASE_COMMIT}"
|
||||
echo "IMAGE_TAG: ${IMAGE_TAG}"
|
||||
echo "IMAGE_TAG_LATEST: ${IMAGE_TAG_LATEST}"
|
||||
|
||||
# print build configuration
|
||||
echo "--- :mag: Build configuration"
|
||||
echo "TARGET: ${TARGET}"
|
||||
echo "CI HCL URL: ${CI_HCL_URL}"
|
||||
echo "vLLM bake file: ${VLLM_BAKE_FILE}"
|
||||
echo "BUILDER_NAME: ${BUILDER_NAME}"
|
||||
echo "CI_HCL_PATH: ${CI_HCL_PATH}"
|
||||
echo "BUILDKIT_SOCKET: ${BUILDKIT_SOCKET}"
|
||||
|
||||
echo "--- :mag: Cache tags"
|
||||
echo "CACHE_TO: ${CACHE_TO}"
|
||||
echo "CACHE_FROM: ${CACHE_FROM}"
|
||||
echo "CACHE_FROM_BASE_BRANCH: ${CACHE_FROM_BASE_BRANCH}"
|
||||
echo "CACHE_FROM_MAIN: ${CACHE_FROM_MAIN}"
|
||||
|
||||
check_and_skip_if_image_exists
|
||||
|
||||
echo "--- :docker: Setting up Docker buildx bake"
|
||||
echo "Target: ${TARGET}"
|
||||
echo "CI HCL URL: ${CI_HCL_URL}"
|
||||
echo "vLLM bake file: ${VLLM_BAKE_FILE}"
|
||||
|
||||
if [[ ! -f "${VLLM_BAKE_FILE}" ]]; then
|
||||
echo "Error: vLLM bake file not found at ${VLLM_BAKE_FILE}"
|
||||
echo "Make sure you're running from the vLLM repository root"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ "${VLLM_USE_PRECOMPILED:-0}" == "1" ]]; then
|
||||
merge_base_commit_build_args="--build-arg VLLM_MERGE_BASE_COMMIT=${VLLM_MERGE_BASE_COMMIT}"
|
||||
else
|
||||
merge_base_commit_build_args=""
|
||||
fi
|
||||
echo "--- :arrow_down: Downloading ci.hcl"
|
||||
curl -sSfL -o "${CI_HCL_PATH}" "${CI_HCL_URL}"
|
||||
echo "Downloaded to ${CI_HCL_PATH}"
|
||||
|
||||
# build
|
||||
docker buildx build --file docker/Dockerfile \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg buildkite_commit=$BUILDKITE_COMMIT \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg TORCH_CUDA_ARCH_LIST="8.0 8.9 9.0 10.0" \
|
||||
--build-arg FI_TORCH_CUDA_ARCH_LIST="8.0 8.9 9.0a 10.0a" \
|
||||
--build-arg VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED:-0}" \
|
||||
${merge_base_commit_build_args} \
|
||||
--cache-from type=registry,ref=${CACHE_FROM},mode=max \
|
||||
--cache-to type=registry,ref=${CACHE_TO},mode=max \
|
||||
--tag ${REGISTRY}/${REPO}:${BUILDKITE_COMMIT} \
|
||||
$( [[ "${BRANCH}" == "main" ]] && echo "--tag ${REGISTRY}/${REPO}:latest" ) \
|
||||
--push \
|
||||
--target test \
|
||||
--progress plain .
|
||||
setup_buildx_builder
|
||||
|
||||
# Compute parent commit for cache fallback (if not already set)
|
||||
resolve_parent_commit
|
||||
export PARENT_COMMIT
|
||||
|
||||
print_bake_config
|
||||
|
||||
echo "--- :docker: Building ${TARGET}"
|
||||
docker --debug buildx bake -f "${VLLM_BAKE_FILE}" -f "${CI_HCL_PATH}" --progress plain "${TARGET}"
|
||||
|
||||
echo "--- :white_check_mark: Build complete"
|
||||
|
||||
@@ -4,7 +4,8 @@ steps:
|
||||
key: image-build
|
||||
depends_on: []
|
||||
commands:
|
||||
- .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $CACHE_FROM $CACHE_TO
|
||||
- if [[ "$BUILDKITE_BRANCH" != "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG; fi
|
||||
- if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG_LATEST; fi
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: -1 # Agent was lost
|
||||
|
||||
@@ -274,14 +274,14 @@ steps:
|
||||
- input-release-version
|
||||
- build-wheels
|
||||
|
||||
- label: "Upload release wheels to PyPI"
|
||||
- label: "Upload release wheels to PyPI and GitHub"
|
||||
depends_on:
|
||||
- block-upload-release-wheels
|
||||
id: upload-release-wheels
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
commands:
|
||||
- "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
|
||||
- "bash .buildkite/scripts/upload-release-wheels.sh"
|
||||
|
||||
# =============================================================================
|
||||
# ROCm Release Pipeline (x86_64 only)
|
||||
@@ -638,93 +638,9 @@ steps:
|
||||
depends_on:
|
||||
- step: upload-rocm-wheels
|
||||
allow_failure: true
|
||||
- step: input-release-version
|
||||
allow_failure: true
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-rocm-release.sh"
|
||||
env:
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
|
||||
# ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
|
||||
# This is the job to create https://wheels.vllm.ai/rocm/ index allowing
|
||||
# users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
|
||||
- block: "Generate Root Index for ROCm Wheels for Release"
|
||||
key: block-generate-root-index-rocm-wheels
|
||||
depends_on: upload-rocm-wheels
|
||||
|
||||
- label: ":package: Generate Root Index for ROCm Wheels for Release"
|
||||
depends_on: block-generate-root-index-rocm-wheels
|
||||
id: generate-root-index-rocm-wheels
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
commands:
|
||||
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
|
||||
env:
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
VARIANT: "rocm700"
|
||||
|
||||
# ROCm Job 5: Build ROCm Release Docker Image
|
||||
- label: ":rocm: :docker: Build ROCm Release Docker Image"
|
||||
id: build-rocm-release-image
|
||||
depends_on:
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- |
|
||||
set -euo pipefail
|
||||
|
||||
# Login to ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
# Download Docker image from S3 (set by build-rocm-base-wheels)
|
||||
DOCKER_IMAGE_S3_PATH="$$(buildkite-agent meta-data get rocm-docker-image-s3-path 2>/dev/null || echo '')"
|
||||
if [ -z "$${DOCKER_IMAGE_S3_PATH}" ]; then
|
||||
echo "ERROR: rocm-docker-image-s3-path metadata not found"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Downloading base image from $${DOCKER_IMAGE_S3_PATH}"
|
||||
mkdir -p artifacts/rocm-docker-image
|
||||
aws s3 cp "$${DOCKER_IMAGE_S3_PATH}" artifacts/rocm-docker-image/rocm-base-image.tar.gz
|
||||
|
||||
# Load base Docker image
|
||||
echo "Loading base Docker image..."
|
||||
LOAD_OUTPUT=$$(gunzip -c artifacts/rocm-docker-image/rocm-base-image.tar.gz | docker load)
|
||||
BASE_IMAGE_TAG=$$(echo "$${LOAD_OUTPUT}" | grep "Loaded image:" | sed 's/Loaded image: //')
|
||||
echo "Loaded base image: $${BASE_IMAGE_TAG}"
|
||||
|
||||
# Tag and push the base image to ECR
|
||||
docker tag "$${BASE_IMAGE_TAG}" public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base
|
||||
docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base
|
||||
echo "Pushed base image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm-base"
|
||||
|
||||
# Get GPU architectures from meta-data
|
||||
PYTORCH_ROCM_ARCH="$$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo '')"
|
||||
PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH:-gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151}"
|
||||
|
||||
# Build vLLM ROCm release image using cached base
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg BASE_IMAGE="$${BASE_IMAGE_TAG}" \
|
||||
--build-arg ARG_PYTORCH_ROCM_ARCH="$${PYTORCH_ROCM_ARCH}" \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
|
||||
--build-arg SCCACHE_REGION_NAME=us-west-2 \
|
||||
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile.rocm .
|
||||
|
||||
# Push to ECR
|
||||
docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
|
||||
echo "Pushed: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
|
||||
@@ -11,80 +11,51 @@ fi
|
||||
buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
|
||||
To download the wheel (by commit):
|
||||
\`\`\`
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_aarch64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux1_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux2014_aarch64.whl .
|
||||
|
||||
(Optional) For CUDA 13.0:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
|
||||
(Optional) For CPU:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
|
||||
\`\`\`
|
||||
|
||||
To download the wheel (by version):
|
||||
\`\`\`
|
||||
aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux1_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux2014_aarch64.whl .
|
||||
|
||||
aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu129/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux1_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${RELEASE_VERSION}+cu130/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux1_x86_64.whl .
|
||||
\`\`\`
|
||||
|
||||
To download and upload the image:
|
||||
|
||||
\`\`\`
|
||||
Download images:
|
||||
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
|
||||
|
||||
Tag and push images:
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
|
||||
docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64
|
||||
docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
docker push vllm/vllm-openai:latest-x86_64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 vllm/vllm-openai:x86_64-cu130
|
||||
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:latest-x86_64-cu130
|
||||
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
|
||||
docker push vllm/vllm-openai:latest-x86_64-cu130
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
|
||||
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
|
||||
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker push vllm/vllm-openai:latest-aarch64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 vllm/vllm-openai:aarch64-cu130
|
||||
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker push vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:latest
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION}-rocm
|
||||
docker push vllm/vllm-openai-rocm:latest
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-rocm
|
||||
|
||||
Create multi-arch manifest:
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
docker push vllm/vllm-openai-rocm:latest-base
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai:rocm
|
||||
docker tag vllm/vllm-openai:rocm vllm/vllm-openai:latest-rocm
|
||||
docker tag vllm/vllm-openai:rocm vllm/vllm-openai:v${RELEASE_VERSION}-rocm
|
||||
docker push vllm/vllm-openai:latest-rocm
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-rocm
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest
|
||||
docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker manifest push vllm/vllm-openai:latest
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest-cu130
|
||||
docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker manifest push vllm/vllm-openai:latest-cu130
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
|
||||
\`\`\`
|
||||
EOF
|
||||
|
||||
@@ -3,32 +3,25 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Generate Buildkite annotation for ROCm wheel release
|
||||
|
||||
set -ex
|
||||
|
||||
# Get build configuration from meta-data
|
||||
# Extract ROCm version dynamically from Dockerfile.rocm_base
|
||||
# BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.0-complete -> extracts "7.0"
|
||||
# BASE_IMAGE format: rocm/dev-ubuntu-22.04:7.1-complete -> extracts "7.1"
|
||||
ROCM_VERSION=$(grep -E '^ARG BASE_IMAGE=' docker/Dockerfile.rocm_base | sed -E 's/.*:([0-9]+\.[0-9]+).*/\1/' || echo "unknown")
|
||||
PYTHON_VERSION=$(buildkite-agent meta-data get rocm-python-version 2>/dev/null || echo "3.12")
|
||||
PYTORCH_ROCM_ARCH=$(buildkite-agent meta-data get rocm-pytorch-rocm-arch 2>/dev/null || echo "gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1200;gfx1201;gfx1150;gfx1151")
|
||||
|
||||
# TODO: Enable the nightly build for ROCm
|
||||
# Get release version, default to 1.0.0.dev for nightly/per-commit builds
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null || echo "")
|
||||
if [ -z "${RELEASE_VERSION}" ]; then
|
||||
RELEASE_VERSION="1.0.0.dev"
|
||||
fi
|
||||
|
||||
# S3 URLs
|
||||
S3_BUCKET="${S3_BUCKET:-vllm-wheels}"
|
||||
S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}"
|
||||
S3_URL="http://${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com"
|
||||
S3_URL="https://${S3_BUCKET}.s3.${S3_REGION}.amazonaws.com"
|
||||
ROCM_PATH="rocm/${BUILDKITE_COMMIT}"
|
||||
|
||||
# Format ROCm version for path (e.g., "7.1" -> "rocm710")
|
||||
ROCM_VERSION_PATH="rocm$(echo ${ROCM_VERSION} | tr -d '.')"
|
||||
ROCM_PATH="rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}"
|
||||
buildkite-agent annotate --style 'success' --context 'rocm-release-workflow' << EOF
|
||||
## ROCm Wheel and Docker Image Releases
|
||||
## :rocm: ROCm Wheel Release
|
||||
|
||||
### Build Configuration
|
||||
| Setting | Value |
|
||||
|---------|-------|
|
||||
@@ -41,72 +34,41 @@ buildkite-agent annotate --style 'success' --context 'rocm-release-workflow' <<
|
||||
### :package: Installation
|
||||
|
||||
**Install from this build (by commit):**
|
||||
|
||||
\`\`\`bash
|
||||
pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com
|
||||
uv pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/{rocm_variant}/
|
||||
|
||||
# Example for ROCm ${ROCM_VERSION}:
|
||||
pip install vllm --extra-index-url ${S3_URL}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com
|
||||
# Example:
|
||||
uv pip install vllm --extra-index-url ${S3_URL}/${ROCM_PATH}/rocm700/
|
||||
\`\`\`
|
||||
|
||||
**Install from nightly (if published):**
|
||||
|
||||
\`\`\`bash
|
||||
pip install vllm --extra-index-url ${S3_URL}/rocm/nightly/ --trusted-host ${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com
|
||||
uv pip install vllm --extra-index-url ${S3_URL}/rocm/nightly/
|
||||
\`\`\`
|
||||
|
||||
### :floppy_disk: Download Wheels Directly
|
||||
|
||||
\`\`\`bash
|
||||
# List all ROCm wheels
|
||||
aws s3 ls s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/
|
||||
aws s3 ls s3://${S3_BUCKET}/${ROCM_PATH}/
|
||||
|
||||
# Download specific wheels
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/vllm-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torch-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/triton-kernels-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchvision-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/torchaudio-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/amdsmi-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/aiter-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}/flash-attn-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/vllm-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/torch-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/triton_rocm-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/torchvision-*.whl .
|
||||
aws s3 cp s3://${S3_BUCKET}/${ROCM_PATH}/amdsmi-*.whl .
|
||||
\`\`\`
|
||||
|
||||
### :gear: Included Packages
|
||||
- **vllm**: vLLM with ROCm support
|
||||
- **torch**: PyTorch built for ROCm ${ROCM_VERSION}
|
||||
- **triton**: Triton
|
||||
- **triton-kernels**: Triton kernels
|
||||
- **triton_rocm**: Triton built for ROCm
|
||||
- **torchvision**: TorchVision for ROCm PyTorch
|
||||
- **torchaudio**: Torchaudio for ROCm PyTorch
|
||||
- **amdsmi**: AMD SMI Python bindings
|
||||
- **aiter**: Aiter for ROCm
|
||||
- **flash-attn**: Flash Attention for ROCm
|
||||
|
||||
### :warning: Notes
|
||||
- These wheels are built for **ROCm ${ROCM_VERSION}** and will NOT work with CUDA GPUs
|
||||
- Supported GPU architectures: ${PYTORCH_ROCM_ARCH}
|
||||
- Platform: Linux x86_64 only
|
||||
|
||||
### :package: Docker Image Release
|
||||
|
||||
To download and upload the image:
|
||||
|
||||
\`\`\`
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
docker push vllm/vllm-openai-rocm:latest-base
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
docker push vllm/vllm-openai-rocm:latest
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
\`\`\`
|
||||
|
||||
EOF
|
||||
|
||||
+44
-10
@@ -7,19 +7,17 @@ SUBPATH=$BUILDKITE_COMMIT
|
||||
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
|
||||
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version)
|
||||
GIT_VERSION=$(git describe --exact-match --tags $BUILDKITE_COMMIT 2>/dev/null)
|
||||
|
||||
echo "Release version from Buildkite: $RELEASE_VERSION"
|
||||
|
||||
if [[ -z "$GIT_VERSION" ]]; then
|
||||
GIT_VERSION=$(git describe --exact-match --tags $BUILDKITE_COMMIT 2>/dev/null)
|
||||
if [ -z "$GIT_VERSION" ]; then
|
||||
echo "[FATAL] Not on a git tag, cannot create release."
|
||||
exit 1
|
||||
else
|
||||
echo "Git version for commit $BUILDKITE_COMMIT: $GIT_VERSION"
|
||||
fi
|
||||
# sanity check for version mismatch
|
||||
if [[ "$RELEASE_VERSION" != "$GIT_VERSION" ]]; then
|
||||
if [[ "$FORCE_RELEASE_IGNORE_VERSION_MISMATCH" == "true" ]]; then
|
||||
if [ "$RELEASE_VERSION" != "$GIT_VERSION" ]; then
|
||||
if [ "$FORCE_RELEASE_IGNORE_VERSION_MISMATCH" == "true" ]; then
|
||||
echo "[WARNING] Force release and ignore version mismatch"
|
||||
else
|
||||
echo "[FATAL] Release version from Buildkite does not match Git version."
|
||||
@@ -29,7 +27,7 @@ fi
|
||||
PURE_VERSION=${RELEASE_VERSION#v} # remove leading 'v'
|
||||
|
||||
# check pypi token
|
||||
if [[ -z "$PYPI_TOKEN" ]]; then
|
||||
if [ -z "$PYPI_TOKEN" ]; then
|
||||
echo "[FATAL] PYPI_TOKEN is not set."
|
||||
exit 1
|
||||
else
|
||||
@@ -37,8 +35,41 @@ else
|
||||
export TWINE_PASSWORD="$PYPI_TOKEN"
|
||||
fi
|
||||
|
||||
# check github token
|
||||
if [ -z "$GITHUB_TOKEN" ]; then
|
||||
echo "[FATAL] GITHUB_TOKEN is not set."
|
||||
exit 1
|
||||
else
|
||||
export GH_TOKEN="$GITHUB_TOKEN"
|
||||
fi
|
||||
|
||||
set -x # avoid printing secrets above
|
||||
|
||||
# download gh CLI from github
|
||||
# Get latest gh CLI version from GitHub API
|
||||
GH_VERSION=$(curl -s https://api.github.com/repos/cli/cli/releases/latest | grep '"tag_name":' | sed -E 's/.*"([^"]+)".*/\1/' | sed 's/^v//')
|
||||
if [ -z "$GH_VERSION" ]; then
|
||||
echo "[FATAL] Failed to get latest gh CLI version from GitHub"
|
||||
exit 1
|
||||
fi
|
||||
echo "Downloading gh CLI version: $GH_VERSION"
|
||||
GH_TARBALL="gh_${GH_VERSION}_linux_amd64.tar.gz"
|
||||
GH_URL="https://github.com/cli/cli/releases/download/v${GH_VERSION}/${GH_TARBALL}"
|
||||
GH_INSTALL_DIR="/tmp/gh-install"
|
||||
mkdir -p "$GH_INSTALL_DIR"
|
||||
pushd "$GH_INSTALL_DIR"
|
||||
curl -L -o "$GH_TARBALL" "$GH_URL"
|
||||
tar -xzf "$GH_TARBALL"
|
||||
GH_BIN=$(realpath $(find . -name "gh" -type f -executable | head -n 1))
|
||||
if [ -z "$GH_BIN" ]; then
|
||||
echo "[FATAL] Failed to find gh CLI executable"
|
||||
exit 1
|
||||
fi
|
||||
echo "gh CLI downloaded successfully, version: $($GH_BIN --version)"
|
||||
echo "Last 5 releases on GitHub:" # as a sanity check of gh and GH_TOKEN
|
||||
command "$GH_BIN" release list --limit 5
|
||||
popd
|
||||
|
||||
# install twine from pypi
|
||||
python3 -m venv /tmp/vllm-release-env
|
||||
source /tmp/vllm-release-env/bin/activate
|
||||
@@ -58,13 +89,16 @@ echo "Wheels copied to local directory"
|
||||
git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" $BUILDKITE_COMMIT
|
||||
ls -la $DIST_DIR
|
||||
|
||||
|
||||
# upload wheels to PyPI (only default variant, i.e. files without '+' in the name)
|
||||
PYPI_WHEEL_FILES=$(find $DIST_DIR -name "vllm-${PURE_VERSION}*.whl" -not -name "*+*")
|
||||
if [[ -z "$PYPI_WHEEL_FILES" ]]; then
|
||||
if [ -z "$PYPI_WHEEL_FILES" ]; then
|
||||
echo "No default variant wheels found, quitting..."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
python3 -m twine check $PYPI_WHEEL_FILES
|
||||
python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
|
||||
python3 -m twine --non-interactive --verbose upload $PYPI_WHEEL_FILES
|
||||
echo "Wheels uploaded to PyPI"
|
||||
|
||||
# create release on GitHub with the release version and all wheels
|
||||
command "$GH_BIN" release create $GIT_VERSION -d --latest --notes-from-tag --verify-tag $DIST_DIR/*.whl
|
||||
@@ -1131,7 +1131,7 @@ steps:
|
||||
- csrc/quantization/cutlass_w8a8/moe/
|
||||
- vllm/model_executor/layers/fused_moe/cutlass_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
|
||||
- vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py
|
||||
- vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
|
||||
- vllm/v1/attention/backends/flashinfer.py
|
||||
- vllm/v1/attention/backends/mla/cutlass_mla.py
|
||||
|
||||
@@ -1017,7 +1017,7 @@ steps:
|
||||
- csrc/quantization/cutlass_w8a8/moe/
|
||||
- vllm/model_executor/layers/fused_moe/cutlass_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
|
||||
- vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py
|
||||
- vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
|
||||
- vllm/v1/attention/backends/flashinfer.py
|
||||
- vllm/v1/attention/backends/mla/cutlass_mla.py
|
||||
@@ -1316,7 +1316,7 @@ steps:
|
||||
- pytest -v -s distributed/test_distributed_oot.py
|
||||
- pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s models/test_oot_registration.py # it needs a clean process
|
||||
- pytest -v -s plugins/lora_resolvers # unit tests for in-tree lora resolver plugins
|
||||
- pytest -v -s plugins/lora_resolvers # unit tests for lora resolver plugins
|
||||
|
||||
- label: Pipeline + Context Parallelism Test # 45min
|
||||
timeout_in_minutes: 60
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: V1 attention (H100)
|
||||
timeout_in_minutes: 30
|
||||
gpu: h100
|
||||
device: h100
|
||||
source_file_dependencies:
|
||||
- vllm/config/attention.py
|
||||
- vllm/model_executor/layers/attention
|
||||
@@ -15,7 +15,7 @@ steps:
|
||||
|
||||
- label: V1 attention (B200)
|
||||
timeout_in_minutes: 30
|
||||
gpu: b200
|
||||
device: b200
|
||||
source_file_dependencies:
|
||||
- vllm/config/attention.py
|
||||
- vllm/model_executor/layers/attention
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: Fusion and Compile Tests (B200)
|
||||
timeout_in_minutes: 40
|
||||
working_dir: "/vllm-workspace/"
|
||||
gpu: b200
|
||||
device: b200
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/fp4/
|
||||
- vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
|
||||
@@ -26,7 +26,7 @@ steps:
|
||||
- nvidia-smi
|
||||
- pytest -v -s tests/compile/test_fusion_attn.py
|
||||
- pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
|
||||
# this runner has 2 GPUs available even though num_gpus=2 is not set
|
||||
# this runner has 2 GPUs available even though num_devices=2 is not set
|
||||
- pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
|
||||
# Limit to Inductor partition, no custom ops, and allreduce & attn fusion to reduce running time
|
||||
# Wrap with quotes to escape yaml
|
||||
@@ -37,9 +37,9 @@ steps:
|
||||
- label: Fusion E2E (2 GPUs)(B200)
|
||||
timeout_in_minutes: 40
|
||||
working_dir: "/vllm-workspace/"
|
||||
gpu: b200
|
||||
device: b200
|
||||
optional: true
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/fp4/
|
||||
- vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: Distributed Comm Ops
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed
|
||||
- tests/distributed
|
||||
@@ -18,7 +18,7 @@ steps:
|
||||
- label: Distributed (2 GPUs)
|
||||
timeout_in_minutes: 90
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/compilation/
|
||||
- vllm/distributed/
|
||||
@@ -54,7 +54,7 @@ steps:
|
||||
- label: Distributed Tests (4 GPUs)
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 4
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/
|
||||
- tests/distributed/test_utils
|
||||
@@ -103,8 +103,8 @@ steps:
|
||||
|
||||
- label: Distributed Tests (8 GPUs)(H100)
|
||||
timeout_in_minutes: 10
|
||||
gpu: h100
|
||||
num_gpus: 8
|
||||
device: h100
|
||||
num_devices: 8
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- examples/offline_inference/torchrun_dp_example.py
|
||||
@@ -120,9 +120,9 @@ steps:
|
||||
- torchrun --nproc-per-node=8 ../examples/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
|
||||
|
||||
- label: Distributed Tests (4 GPUs)(A100)
|
||||
gpu: a100
|
||||
device: a100
|
||||
optional: true
|
||||
num_gpus: 4
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
commands:
|
||||
@@ -133,26 +133,34 @@ steps:
|
||||
- TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest -v -s -x lora/test_mixtral.py
|
||||
|
||||
- label: Distributed Tests (2 GPUs)(H200)
|
||||
gpu: h200
|
||||
- label: Sequence Parallel Tests (H100)
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
# Run sequence parallel tests
|
||||
- pytest -v -s tests/distributed/test_sequence_parallel.py
|
||||
- pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
|
||||
|
||||
- label: Distributed Tests (2 GPUs)(H100)
|
||||
device: h100
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/"
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
commands:
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_async_tp.py
|
||||
- pytest -v -s tests/compile/distributed/test_sequence_parallelism.py
|
||||
- pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/distributed/test_sequence_parallel.py
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- CUDA_VISIBLE_DEVICES=1,2 VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
|
||||
- label: Distributed Tests (2 GPUs)(B200)
|
||||
gpu: b200
|
||||
device: b200
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/"
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
commands:
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
|
||||
@@ -161,8 +169,9 @@ steps:
|
||||
- label: 2 Node Test (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
num_nodes: 2
|
||||
no_plugin: true
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
@@ -176,7 +185,7 @@ steps:
|
||||
- label: Distributed NixlConnector PD accuracy (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 4
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
@@ -184,10 +193,21 @@ steps:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Pipeline + Context Parallelism (4 GPUs))
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 4
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
@@ -196,4 +216,46 @@ steps:
|
||||
- tests/distributed/
|
||||
commands:
|
||||
- pytest -v -s distributed/test_pp_cudagraph.py
|
||||
- pytest -v -s distributed/test_pipeline_parallel.py
|
||||
- pytest -v -s distributed/test_pipeline_parallel.py
|
||||
|
||||
- label: Hopper Fusion E2E Tests (H100)
|
||||
timeout_in_minutes: 70
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/fp4/
|
||||
- vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
|
||||
- vllm/v1/attention/backends/flashinfer.py
|
||||
- vllm/compilation/
|
||||
# can affect pattern matching
|
||||
- vllm/model_executor/layers/layernorm.py
|
||||
- vllm/model_executor/layers/activation.py
|
||||
- vllm/model_executor/layers/quantization/input_quant_fp8.py
|
||||
- tests/compile/test_fusion_attn.py
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
# skip Llama-4 since it does not fit on this device
|
||||
- pytest -v -s tests/compile/test_fusion_attn.py -k 'not Llama-4'
|
||||
|
||||
- label: Hopper Fusion Distributed E2E Tests (2xH100)
|
||||
timeout_in_minutes: 70
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/fp4/
|
||||
- vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
|
||||
- vllm/v1/attention/backends/flashinfer.py
|
||||
- vllm/compilation/
|
||||
# can affect pattern matching
|
||||
- vllm/model_executor/layers/layernorm.py
|
||||
- vllm/model_executor/layers/activation.py
|
||||
- vllm/model_executor/layers/quantization/input_quant_fp8.py
|
||||
- tests/compile/distributed/test_fusions_e2e.py
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
# Run all e2e fusion tests
|
||||
- pytest -v -s tests/compile/distributed/test_fusions_e2e.py -k 'not Llama-4'
|
||||
- pytest -v -s tests/compile/distributed/test_fusion_all_reduce.py
|
||||
|
||||
@@ -4,27 +4,27 @@ depends_on:
|
||||
steps:
|
||||
- label: DeepSeek V2-Lite Accuracy
|
||||
timeout_in_minutes: 60
|
||||
gpu: h100
|
||||
device: h100
|
||||
optional: true
|
||||
num_gpus: 4
|
||||
num_devices: 4
|
||||
working_dir: "/vllm-workspace"
|
||||
commands:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010
|
||||
|
||||
- label: Qwen3-30B-A3B-FP8-block Accuracy
|
||||
timeout_in_minutes: 60
|
||||
gpu: h100
|
||||
device: h100
|
||||
optional: true
|
||||
num_gpus: 4
|
||||
num_devices: 4
|
||||
working_dir: "/vllm-workspace"
|
||||
commands:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
|
||||
|
||||
- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
|
||||
timeout_in_minutes: 60
|
||||
gpu: b200
|
||||
device: b200
|
||||
optional: true
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
working_dir: "/vllm-workspace"
|
||||
commands:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
|
||||
@@ -33,10 +33,11 @@ steps:
|
||||
timeout_in_minutes: 30
|
||||
optional: true
|
||||
soft_fail: true
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
working_dir: "/vllm-workspace"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- .buildkite/scripts/run-prime-rl-test.sh
|
||||
commands:
|
||||
- nvidia-smi
|
||||
- bash .buildkite/scripts/run-prime-rl-test.sh
|
||||
|
||||
@@ -23,4 +23,8 @@ steps:
|
||||
# TODO: accuracy does not match, whether setting
|
||||
# VLLM_USE_FLASHINFER_SAMPLER or not on H100.
|
||||
- pytest -v -s v1/e2e
|
||||
- pytest -v -s v1/engine
|
||||
# Run this test standalone for now;
|
||||
# need to untangle use (implicit) use of spawn/fork across the tests.
|
||||
- pytest -v -s v1/engine/test_preprocess_error_handling.py
|
||||
# Run the rest of v1/engine tests
|
||||
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
|
||||
|
||||
@@ -14,7 +14,7 @@ steps:
|
||||
- label: EPLB Execution
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 4
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/eplb
|
||||
- tests/distributed/test_eplb_execute.py
|
||||
|
||||
@@ -57,8 +57,8 @@ steps:
|
||||
|
||||
- label: Kernels DeepGEMM Test (H100)
|
||||
timeout_in_minutes: 45
|
||||
gpu: h100
|
||||
num_gpus: 1
|
||||
device: h100
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- tools/install_deepgemm.sh
|
||||
- vllm/utils/deep_gemm.py
|
||||
@@ -77,7 +77,7 @@ steps:
|
||||
- label: Kernels (B200)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/"
|
||||
gpu: b200
|
||||
device: b200
|
||||
# optional: true
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/fp4/
|
||||
@@ -85,7 +85,7 @@ steps:
|
||||
- csrc/quantization/cutlass_w8a8/moe/
|
||||
- vllm/model_executor/layers/fused_moe/cutlass_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
|
||||
- vllm/model_executor/layers/fused_moe/flashinfer_a2a_prepare_finalize.py
|
||||
- vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
|
||||
- vllm/v1/attention/backends/flashinfer.py
|
||||
- vllm/v1/attention/backends/mla/cutlass_mla.py
|
||||
@@ -114,4 +114,55 @@ steps:
|
||||
- pytest -v -s tests/kernels/moe/test_nvfp4_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_flashinfer.py
|
||||
- pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
|
||||
# e2e
|
||||
- pytest -v -s tests/models/quantization/test_nvfp4.py
|
||||
|
||||
- label: Kernels Helion Test
|
||||
timeout_in_minutes: 30
|
||||
device: h100
|
||||
source_file_dependencies:
|
||||
- vllm/utils/import_utils.py
|
||||
- tests/kernels/helion/
|
||||
commands:
|
||||
- pip install helion
|
||||
- pytest -v -s kernels/helion/
|
||||
|
||||
|
||||
- label: Kernels FP8 MoE Test (1 H100)
|
||||
timeout_in_minutes: 90
|
||||
device: h100
|
||||
num_devices: 1
|
||||
optional: true
|
||||
commands:
|
||||
- pytest -v -s kernels/moe/test_cutlass_moe.py
|
||||
- pytest -v -s kernels/moe/test_flashinfer.py
|
||||
- pytest -v -s kernels/moe/test_gpt_oss_triton_kernels.py
|
||||
- pytest -v -s kernels/moe/test_modular_oai_triton_moe.py
|
||||
- pytest -v -s kernels/moe/test_moe.py
|
||||
# - pytest -v -s kernels/moe/test_block_fp8.py - failing on main
|
||||
- pytest -v -s kernels/moe/test_block_int8.py
|
||||
- pytest -v -s kernels/moe/test_triton_moe_no_act_mul.py
|
||||
- pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py
|
||||
|
||||
- label: Kernels FP8 MoE Test (2 H100s)
|
||||
timeout_in_minutes: 90
|
||||
device: h100
|
||||
num_devices: 2
|
||||
optional: true
|
||||
commands:
|
||||
- pytest -v -s kernels/moe/test_deepep_deepgemm_moe.py
|
||||
- pytest -v -s kernels/moe/test_deepep_moe.py
|
||||
- pytest -v -s kernels/moe/test_pplx_cutlass_moe.py
|
||||
# - pytest -v -s kernels/moe/test_pplx_moe.py - failing on main
|
||||
|
||||
- label: Kernels Fp4 MoE Test (B200)
|
||||
timeout_in_minutes: 60
|
||||
device: b200
|
||||
num_devices: 1
|
||||
optional: true
|
||||
commands:
|
||||
- pytest -v -s kernels/moe/test_cutedsl_moe.py
|
||||
- pytest -v -s kernels/moe/test_flashinfer_moe.py
|
||||
- pytest -v -s kernels/moe/test_nvfp4_moe.py
|
||||
- pytest -v -s kernels/moe/test_ocp_mx_moe.py
|
||||
|
||||
@@ -12,9 +12,9 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt
|
||||
|
||||
- label: LM Eval Large Models (4 GPUs)(A100)
|
||||
gpu: a100
|
||||
device: a100
|
||||
optional: true
|
||||
num_gpus: 4
|
||||
num_devices: 4
|
||||
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
@@ -24,9 +24,9 @@ steps:
|
||||
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
|
||||
|
||||
- label: LM Eval Large Models (4 GPUs)(H100)
|
||||
gpu: h100
|
||||
device: h100
|
||||
optional: true
|
||||
num_gpus: 4
|
||||
num_devices: 4
|
||||
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
@@ -37,10 +37,39 @@ steps:
|
||||
|
||||
- label: LM Eval Small Models (B200)
|
||||
timeout_in_minutes: 120
|
||||
gpu: b200
|
||||
device: b200
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
|
||||
|
||||
- label: LM Eval Large Models (H200)
|
||||
timeout_in_minutes: 60
|
||||
device: h200
|
||||
optional: true
|
||||
num_devices: 8
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
|
||||
|
||||
- label: MoE Refactor Integration Test (H100 - TEMPORARY)
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100.txt
|
||||
|
||||
- label: MoE Refactor Integration Test (B200 - TEMPORARY)
|
||||
gpu: b200
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt
|
||||
|
||||
- label: MoE Refactor Integration Test (B200 DP - TEMPORARY)
|
||||
device: b200
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
|
||||
|
||||
@@ -14,7 +14,7 @@ steps:
|
||||
|
||||
- label: LoRA TP (Distributed)
|
||||
timeout_in_minutes: 30
|
||||
num_gpus: 4
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/lora
|
||||
- tests/lora
|
||||
|
||||
@@ -31,7 +31,7 @@ steps:
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/v1
|
||||
no_gpu: true
|
||||
device: cpu
|
||||
commands:
|
||||
# split the test to avoid interference
|
||||
- pytest -v -s -m 'cpu_test' v1/core
|
||||
@@ -82,7 +82,7 @@ steps:
|
||||
|
||||
- label: Metrics, Tracing (2 GPUs)
|
||||
timeout_in_minutes: 20
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/v1/tracing
|
||||
@@ -127,7 +127,7 @@ steps:
|
||||
- tests/tool_parsers
|
||||
- tests/transformers_utils
|
||||
- tests/config
|
||||
no_gpu: true
|
||||
device: cpu
|
||||
commands:
|
||||
- python3 standalone_tests/lazy_imports.py
|
||||
- pytest -v -s test_inputs.py
|
||||
@@ -142,7 +142,7 @@ steps:
|
||||
- label: GPT-OSS Eval (B200)
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/"
|
||||
gpu: b200
|
||||
device: b200
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- tests/evals/gpt_oss
|
||||
@@ -155,7 +155,7 @@ steps:
|
||||
|
||||
- label: Batch Invariance (H100)
|
||||
timeout_in_minutes: 25
|
||||
gpu: h100
|
||||
device: h100
|
||||
source_file_dependencies:
|
||||
- vllm/v1/attention
|
||||
- vllm/model_executor/layers
|
||||
|
||||
@@ -44,7 +44,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/test_utils.py
|
||||
- tests/models/test_vision.py
|
||||
no_gpu: true
|
||||
device: cpu
|
||||
commands:
|
||||
- pytest -v -s models/test_utils.py models/test_vision.py
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: Distributed Model Tests (2 GPUs)
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/model_loader/sharded_state_loader.py
|
||||
- vllm/model_executor/models/
|
||||
|
||||
@@ -18,7 +18,7 @@ steps:
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
no_gpu: true
|
||||
device: cpu
|
||||
commands:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: Plugin Tests (2 GPUs)
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/plugins/
|
||||
- tests/plugins/
|
||||
|
||||
@@ -16,14 +16,14 @@ steps:
|
||||
# https://github.com/pytorch/ao/issues/2919, we'll have to skip new torchao tests for now
|
||||
# we can only upgrade after this is resolved
|
||||
# TODO(jerryzh168): resolve the above comment
|
||||
- uv pip install --system torchao==0.13.0 --index-url https://download.pytorch.org/whl/cu129
|
||||
- uv pip install --system torchao==0.14.1 --index-url https://download.pytorch.org/whl/cu129
|
||||
- uv pip install --system conch-triton-kernels
|
||||
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
|
||||
|
||||
- label: Quantized MoE Test (B200)
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/"
|
||||
gpu: b200
|
||||
device: b200
|
||||
source_file_dependencies:
|
||||
- tests/quantization/test_blackwell_moe.py
|
||||
- vllm/model_executor/models/deepseek_v2.py
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: Weight Loading Multiple GPU # 33min
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 2
|
||||
num_devices: 2
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -15,8 +15,8 @@ steps:
|
||||
|
||||
- label: Weight Loading Multiple GPU - Large Models # optional
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_gpus: 2
|
||||
gpu: a100
|
||||
num_devices: 2
|
||||
device: a100
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -197,7 +197,7 @@ def bench_run(
|
||||
)
|
||||
|
||||
kernel = mk.FusedMoEModularKernel(
|
||||
MoEPrepareAndFinalizeNoEP(defer_input_quant=True),
|
||||
MoEPrepareAndFinalizeNoEP(),
|
||||
CutlassExpertsFp4(
|
||||
make_dummy_moe_config(),
|
||||
quant_config=quant_config,
|
||||
@@ -242,7 +242,7 @@ def bench_run(
|
||||
)
|
||||
|
||||
kernel = mk.FusedMoEModularKernel(
|
||||
MoEPrepareAndFinalizeNoEP(defer_input_quant=True),
|
||||
MoEPrepareAndFinalizeNoEP(),
|
||||
CutlassExpertsFp4(
|
||||
make_dummy_moe_config(),
|
||||
quant_config=quant_config,
|
||||
|
||||
@@ -10,8 +10,6 @@ from transformers import AutoConfig
|
||||
|
||||
from vllm.model_executor.layers.fused_moe import fused_topk
|
||||
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
|
||||
_moe_permute,
|
||||
_moe_unpermute_and_reduce,
|
||||
moe_permute,
|
||||
moe_unpermute,
|
||||
)
|
||||
@@ -41,7 +39,6 @@ def benchmark_permute(
|
||||
use_fp8_w8a8: bool,
|
||||
use_int8_w8a16: bool,
|
||||
num_iters: int = 100,
|
||||
use_customized_permute: bool = False,
|
||||
) -> float:
|
||||
# init_dtype = torch.float16 if use_fp8_w8a8 else dtype
|
||||
hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)
|
||||
@@ -64,29 +61,14 @@ def benchmark_permute(
|
||||
input_gating.copy_(gating_output[i])
|
||||
|
||||
def run():
|
||||
if use_customized_permute:
|
||||
(
|
||||
permuted_hidden_states,
|
||||
a1q_scale,
|
||||
first_token_off,
|
||||
inv_perm_idx,
|
||||
m_indices,
|
||||
) = moe_permute(
|
||||
qhidden_states,
|
||||
a1q_scale=None,
|
||||
topk_ids=topk_ids,
|
||||
n_expert=num_experts,
|
||||
expert_map=None,
|
||||
align_block_size=align_block_size,
|
||||
)
|
||||
else:
|
||||
(
|
||||
permuted_hidden_states,
|
||||
a1q_scale,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
inv_perm,
|
||||
) = _moe_permute(qhidden_states, None, topk_ids, num_experts, None, 16)
|
||||
moe_permute(
|
||||
qhidden_states,
|
||||
a1q_scale=None,
|
||||
topk_ids=topk_ids,
|
||||
n_expert=num_experts,
|
||||
expert_map=None,
|
||||
align_block_size=align_block_size,
|
||||
)
|
||||
|
||||
# JIT compilation & warmup
|
||||
run()
|
||||
@@ -131,11 +113,9 @@ def benchmark_unpermute(
|
||||
use_fp8_w8a8: bool,
|
||||
use_int8_w8a16: bool,
|
||||
num_iters: int = 100,
|
||||
use_customized_permute: bool = False,
|
||||
) -> float:
|
||||
# init_dtype = torch.float16 if use_fp8_w8a8 else dtype
|
||||
hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype)
|
||||
output_hidden_states = torch.empty_like(hidden_states)
|
||||
if use_fp8_w8a8:
|
||||
align_block_size = 128 # deepgemm needs 128 m aligned block
|
||||
qhidden_states, scale = _fp8_quantize(hidden_states, None, None)
|
||||
@@ -150,78 +130,37 @@ def benchmark_unpermute(
|
||||
)
|
||||
|
||||
def prepare():
|
||||
if use_customized_permute:
|
||||
(
|
||||
permuted_hidden_states,
|
||||
a1q_scale,
|
||||
first_token_off,
|
||||
inv_perm_idx,
|
||||
m_indices,
|
||||
) = moe_permute(
|
||||
qhidden_states,
|
||||
a1q_scale=None,
|
||||
topk_ids=topk_ids,
|
||||
n_expert=num_experts,
|
||||
expert_map=None,
|
||||
align_block_size=align_block_size,
|
||||
)
|
||||
# convert to fp16/bf16 as gemm output
|
||||
return (
|
||||
permuted_hidden_states.to(dtype),
|
||||
first_token_off,
|
||||
inv_perm_idx,
|
||||
m_indices,
|
||||
)
|
||||
else:
|
||||
(
|
||||
permuted_qhidden_states,
|
||||
a1q_scale,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
inv_perm,
|
||||
) = _moe_permute(
|
||||
qhidden_states, None, topk_ids, num_experts, None, block_m=16
|
||||
)
|
||||
# convert to fp16/bf16 as gemm output
|
||||
return (
|
||||
permuted_qhidden_states.to(dtype),
|
||||
a1q_scale,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
inv_perm,
|
||||
)
|
||||
(
|
||||
permuted_hidden_states,
|
||||
_,
|
||||
first_token_off,
|
||||
inv_perm_idx,
|
||||
_,
|
||||
) = moe_permute(
|
||||
qhidden_states,
|
||||
a1q_scale=None,
|
||||
topk_ids=topk_ids,
|
||||
n_expert=num_experts,
|
||||
expert_map=None,
|
||||
align_block_size=align_block_size,
|
||||
)
|
||||
# convert to fp16/bf16 as gemm output
|
||||
return (
|
||||
permuted_hidden_states.to(dtype),
|
||||
first_token_off,
|
||||
inv_perm_idx,
|
||||
)
|
||||
|
||||
def run(input: tuple):
|
||||
if use_customized_permute:
|
||||
(
|
||||
permuted_hidden_states,
|
||||
first_token_off,
|
||||
inv_perm_idx,
|
||||
m_indices,
|
||||
) = input
|
||||
output = torch.empty_like(hidden_states)
|
||||
moe_unpermute(
|
||||
output,
|
||||
permuted_hidden_states,
|
||||
topk_weights,
|
||||
inv_perm_idx,
|
||||
first_token_off,
|
||||
)
|
||||
else:
|
||||
(
|
||||
permuted_hidden_states,
|
||||
a1q_scale,
|
||||
sorted_token_ids,
|
||||
expert_ids,
|
||||
inv_perm,
|
||||
) = input
|
||||
_moe_unpermute_and_reduce(
|
||||
output_hidden_states,
|
||||
permuted_hidden_states,
|
||||
inv_perm,
|
||||
topk_weights,
|
||||
True,
|
||||
)
|
||||
(permuted_hidden_states, first_token_off, inv_perm_idx) = input
|
||||
output = torch.empty_like(hidden_states)
|
||||
moe_unpermute(
|
||||
output,
|
||||
permuted_hidden_states,
|
||||
topk_weights,
|
||||
inv_perm_idx,
|
||||
first_token_off,
|
||||
)
|
||||
|
||||
# JIT compilation & warmup
|
||||
input = prepare()
|
||||
@@ -276,8 +215,7 @@ class BenchmarkWorker:
|
||||
dtype: torch.dtype,
|
||||
use_fp8_w8a8: bool,
|
||||
use_int8_w8a16: bool,
|
||||
use_customized_permute: bool = False,
|
||||
) -> tuple[dict[str, int], float]:
|
||||
) -> tuple[float, float]:
|
||||
set_random_seed(self.seed)
|
||||
|
||||
permute_time = benchmark_permute(
|
||||
@@ -289,7 +227,6 @@ class BenchmarkWorker:
|
||||
use_fp8_w8a8,
|
||||
use_int8_w8a16,
|
||||
num_iters=100,
|
||||
use_customized_permute=use_customized_permute,
|
||||
)
|
||||
unpermute_time = benchmark_unpermute(
|
||||
num_tokens,
|
||||
@@ -300,7 +237,6 @@ class BenchmarkWorker:
|
||||
use_fp8_w8a8,
|
||||
use_int8_w8a16,
|
||||
num_iters=100,
|
||||
use_customized_permute=use_customized_permute,
|
||||
)
|
||||
return permute_time, unpermute_time
|
||||
|
||||
@@ -347,7 +283,6 @@ def main(args: argparse.Namespace):
|
||||
dtype = torch.float16 if current_platform.is_rocm() else config.dtype
|
||||
use_fp8_w8a8 = args.dtype == "fp8_w8a8"
|
||||
use_int8_w8a16 = args.dtype == "int8_w8a16"
|
||||
use_customized_permute = args.use_customized_permute
|
||||
|
||||
if args.batch_size is None:
|
||||
batch_sizes = [
|
||||
@@ -399,7 +334,6 @@ def main(args: argparse.Namespace):
|
||||
dtype,
|
||||
use_fp8_w8a8,
|
||||
use_int8_w8a16,
|
||||
use_customized_permute,
|
||||
)
|
||||
for batch_size in batch_sizes
|
||||
],
|
||||
@@ -419,7 +353,6 @@ if __name__ == "__main__":
|
||||
parser.add_argument(
|
||||
"--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"
|
||||
)
|
||||
parser.add_argument("--use-customized-permute", action="store_true")
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument("--batch-size", type=int, required=False)
|
||||
parser.add_argument("--trust-remote-code", action="store_true")
|
||||
|
||||
@@ -360,13 +360,14 @@ void onednn_scaled_mm(
|
||||
const std::optional<torch::Tensor>& azp, // [M] or [1]
|
||||
const std::optional<torch::Tensor>& azp_adj, // [M] or [1]
|
||||
const std::optional<torch::Tensor>& bias, // [N]
|
||||
int64_t handler) {
|
||||
const torch::Tensor& handler_tensor) {
|
||||
CPU_KERNEL_GUARD_IN(onednn_scaled_mm)
|
||||
TORCH_CHECK(a.dim() == 2);
|
||||
TORCH_CHECK(a.is_contiguous());
|
||||
TORCH_CHECK(c.is_contiguous());
|
||||
W8A8MatMulPrimitiveHandler* ptr =
|
||||
reinterpret_cast<W8A8MatMulPrimitiveHandler*>(handler);
|
||||
reinterpret_cast<W8A8MatMulPrimitiveHandler*>(
|
||||
handler_tensor.item<int64_t>());
|
||||
const int32_t* azp_ptr = nullptr;
|
||||
if (azp.has_value()) {
|
||||
azp_ptr = azp->data_ptr<int32_t>();
|
||||
@@ -519,13 +520,14 @@ int64_t create_onednn_mm_handler(const torch::Tensor& b,
|
||||
|
||||
void onednn_mm(torch::Tensor& c, // [M, OC], row-major
|
||||
const torch::Tensor& a, // [M, IC], row-major
|
||||
const std::optional<torch::Tensor>& bias, int64_t handler) {
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const torch::Tensor& handler_tensor) {
|
||||
CPU_KERNEL_GUARD_IN(onednn_mm)
|
||||
TORCH_CHECK(a.dim() == 2);
|
||||
TORCH_CHECK(a.stride(-1) == 1);
|
||||
TORCH_CHECK(c.stride(-1) == 1);
|
||||
MatMulPrimitiveHandler* ptr =
|
||||
reinterpret_cast<MatMulPrimitiveHandler*>(handler);
|
||||
reinterpret_cast<MatMulPrimitiveHandler*>(handler_tensor.item<int64_t>());
|
||||
|
||||
// ACL matmuls expect contiguous source tensors
|
||||
#ifdef VLLM_USE_ACL
|
||||
|
||||
@@ -19,13 +19,14 @@ void onednn_scaled_mm(torch::Tensor& c, const torch::Tensor& a,
|
||||
const std::optional<torch::Tensor>& azp,
|
||||
const std::optional<torch::Tensor>& azp_adj,
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
int64_t handler);
|
||||
const torch::Tensor& handler_tensor);
|
||||
|
||||
int64_t create_onednn_mm_handler(const torch::Tensor& b,
|
||||
int64_t primitive_cache_size);
|
||||
|
||||
void onednn_mm(torch::Tensor& c, const torch::Tensor& a,
|
||||
const std::optional<torch::Tensor>& bias, int64_t handler);
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const torch::Tensor& handler_tensor);
|
||||
|
||||
bool is_onednn_acl_supported();
|
||||
|
||||
@@ -196,7 +197,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
// oneDNN GEMM
|
||||
ops.def(
|
||||
"onednn_mm(Tensor! c, Tensor a, Tensor? bias, "
|
||||
"int handler) -> ()");
|
||||
"Tensor handler_tensor) -> ()");
|
||||
ops.impl("onednn_mm", torch::kCPU, &onednn_mm);
|
||||
|
||||
// Check if oneDNN was built with ACL backend
|
||||
@@ -212,7 +213,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
// oneDNN scaled_mm for W8A8 with static per-tensor activation quantization
|
||||
ops.def(
|
||||
"onednn_scaled_mm(Tensor! c, Tensor a, Tensor a_scales, Tensor? azp, "
|
||||
"Tensor? azp_adj, Tensor? bias, int handler) -> ()");
|
||||
"Tensor? azp_adj, Tensor? bias, Tensor handler_tensor) -> ()");
|
||||
ops.impl("onednn_scaled_mm", torch::kCPU, &onednn_scaled_mm);
|
||||
|
||||
// Compute int8 quantized tensor for given scaling factor.
|
||||
|
||||
+1
-14
@@ -227,7 +227,7 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \
|
||||
# This ensures setuptools_scm sees clean repo state for version detection
|
||||
RUN --mount=type=bind,source=.git,target=vllm/.git \
|
||||
cd vllm \
|
||||
&& pip install setuptools_scm regex \
|
||||
&& pip install setuptools_scm \
|
||||
&& VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \
|
||||
&& echo "Detected vLLM version: ${VLLM_VERSION}" \
|
||||
&& echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
|
||||
@@ -342,19 +342,6 @@ RUN mkdir src && mv vllm src/vllm
|
||||
FROM base AS final
|
||||
|
||||
RUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Clean up sccache from release image (not needed at runtime)
|
||||
# This removes the binary and wrappers that may have been installed during build
|
||||
RUN rm -f /usr/bin/sccache || true \
|
||||
&& rm -rf /opt/sccache-wrappers || true
|
||||
|
||||
# Unset sccache environment variables for the release image
|
||||
# This prevents S3 bucket config from leaking into production images
|
||||
ENV SCCACHE_BUCKET=
|
||||
ENV SCCACHE_REGION=
|
||||
ENV SCCACHE_S3_NO_CREDENTIALS=
|
||||
ENV SCCACHE_IDLE_TIMEOUT=
|
||||
|
||||
# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt.
|
||||
# Manually remove it so that later steps of numpy upgrade can continue
|
||||
RUN case "$(which python3)" in \
|
||||
|
||||
@@ -47,6 +47,10 @@ You can tune the performance by adjusting `max_num_batched_tokens`:
|
||||
- For optimal throughput, we recommend setting `max_num_batched_tokens > 8192` especially for smaller models on large GPUs.
|
||||
- If `max_num_batched_tokens` is the same as `max_model_len`, that's almost the equivalent to the V0 default scheduling policy (except that it still prioritizes decodes).
|
||||
|
||||
!!! warning
|
||||
When chunked prefill is disabled, `max_num_batched_tokens` must be greater than `max_model_len`.
|
||||
In that case, if `max_num_batched_tokens < max_model_len`, vLLM may crash at server start‑up.
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
|
||||
@@ -71,7 +71,7 @@ class MyModel(nn.Module):
|
||||
```python
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
input_ids: torch.Tensor | None,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
|
||||
@@ -43,28 +43,73 @@ Further update the model as follows:
|
||||
)
|
||||
```
|
||||
|
||||
- Implement [embed_multimodal][vllm.model_executor.models.interfaces.SupportsMultiModal.embed_multimodal] that returns the embeddings from running the multimodal inputs through the multimodal tokenizer of the model. Below we provide a boilerplate of a typical implementation pattern, but feel free to adjust it to your own needs.
|
||||
- Remove the embedding part from the [forward][torch.nn.Module.forward] method:
|
||||
- Move the multi-modal embedding to [embed_multimodal][vllm.model_executor.models.interfaces.SupportsMultiModal.embed_multimodal].
|
||||
- The text embedding and embedding merge are handled automatically by a default implementation of [embed_input_ids][vllm.model_executor.models.interfaces.SupportsMultiModal.embed_input_ids]. It does not need to be overridden in most cases.
|
||||
|
||||
??? code
|
||||
```diff
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor | None,
|
||||
- pixel_values: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
- if inputs_embeds is None:
|
||||
- inputs_embeds = self.get_input_embeddings()(input_ids)
|
||||
-
|
||||
- if pixel_values is not None:
|
||||
- image_features = self.get_image_features(
|
||||
- pixel_values=pixel_values,
|
||||
- )
|
||||
- special_image_mask = self.get_placeholder_mask(
|
||||
- input_ids,
|
||||
- inputs_embeds=inputs_embeds,
|
||||
- image_features=image_features,
|
||||
- )
|
||||
- inputs_embeds = inputs_embeds.masked_scatter(
|
||||
- special_image_mask,
|
||||
- image_features,
|
||||
- )
|
||||
|
||||
```python
|
||||
def _process_image_input(self, image_input: YourModelImageInputs) -> torch.Tensor:
|
||||
image_features = self.vision_encoder(image_input)
|
||||
return self.multi_modal_projector(image_features)
|
||||
hidden_states = self.language_model(
|
||||
input_ids,
|
||||
positions,
|
||||
intermediate_tensors,
|
||||
inputs_embeds=inputs_embeds,
|
||||
)
|
||||
...
|
||||
|
||||
+ def embed_multimodal(
|
||||
+ self,
|
||||
+ pixel_values: torch.Tensor,
|
||||
+ ) -> MultiModalEmbeddings | None:
|
||||
+ return self.get_image_features(
|
||||
+ pixel_values=pixel_values,
|
||||
+ )
|
||||
```
|
||||
|
||||
def embed_multimodal(
|
||||
self,
|
||||
**kwargs: object,
|
||||
) -> MultiModalEmbeddings | None:
|
||||
# Validate the multimodal input keyword arguments
|
||||
image_input = self._parse_and_validate_image_input(**kwargs)
|
||||
if image_input is None:
|
||||
return None
|
||||
Below we provide a boilerplate of a typical implementation pattern of [embed_multimodal][vllm.model_executor.models.interfaces.SupportsMultiModal.embed_multimodal], but feel free to adjust it to your own needs.
|
||||
|
||||
# Run multimodal inputs through encoder and projector
|
||||
vision_embeddings = self._process_image_input(image_input)
|
||||
return vision_embeddings
|
||||
```
|
||||
```python
|
||||
def _process_image_input(self, image_input: YourModelImageInputs) -> torch.Tensor:
|
||||
image_features = self.vision_encoder(image_input)
|
||||
return self.multi_modal_projector(image_features)
|
||||
|
||||
def embed_multimodal(
|
||||
self,
|
||||
**kwargs: object,
|
||||
) -> MultiModalEmbeddings | None:
|
||||
# Validate the multimodal input keyword arguments
|
||||
image_input = self._parse_and_validate_image_input(**kwargs)
|
||||
if image_input is None:
|
||||
return None
|
||||
|
||||
# Run multimodal inputs through encoder and projector
|
||||
vision_embeddings = self._process_image_input(image_input)
|
||||
return vision_embeddings
|
||||
```
|
||||
|
||||
!!! important
|
||||
The returned `multimodal_embeddings` must be either a **3D [torch.Tensor][]** of shape `(num_items, feature_size, hidden_size)`, or a **list / tuple of 2D [torch.Tensor][]'s** of shape `(feature_size, hidden_size)`, so that `multimodal_embeddings[i]` retrieves the embeddings generated from the `i`-th multimodal data item (e.g, image) of the request.
|
||||
|
||||
@@ -10,7 +10,7 @@ receives a request for a LoRA adapter that hasn't been loaded yet, the resolver
|
||||
to locate and load the adapter from their configured storage locations. This enables:
|
||||
|
||||
- **Dynamic LoRA Loading**: Load adapters on-demand without server restarts
|
||||
- **Multiple Storage Backends**: Support for filesystem, S3, and custom backends. The built-in `lora_filesystem_resolver` requires a local storage path, but custom resolvers can be implemented to fetch from any source.
|
||||
- **Multiple Storage Backends**: Support for filesystem, S3, and custom backends. The built-in `lora_filesystem_resolver` requires a local storage path, while the built-in `hf_hub_resolver` will pull LoRA adapters from Huggingface Hub and proceed in an identical manner. In general, custom resolvers can be implemented to fetch from any source.
|
||||
- **Automatic Discovery**: Seamless integration with existing LoRA workflows
|
||||
- **Scalable Deployment**: Centralized adapter management across multiple vLLM instances
|
||||
|
||||
|
||||
@@ -36,8 +36,7 @@ th {
|
||||
| pplx | batched | fp8,int8 | G,A,T | Y | Y | [`PplxPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.pplx_prepare_finalize.PplxPrepareAndFinalize] |
|
||||
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
|
||||
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
|
||||
| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferAllToAllMoEPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize.FlashInferAllToAllMoEPrepareAndFinalize] |
|
||||
| flashinfer<sup>4</sup> | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferCutlassMoEPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize.FlashInferCutlassMoEPrepareAndFinalize] |
|
||||
| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] |
|
||||
| MoEPrepareAndFinalizeNoEP<sup>5</sup> | standard | fp8,int8 | G,A,T | N | Y | [`MoEPrepareAndFinalizeNoEP`][vllm.model_executor.layers.fused_moe.prepare_finalize.MoEPrepareAndFinalizeNoEP] |
|
||||
| BatchedPrepareAndFinalize<sup>5</sup> | batched | fp8,int8 | G,A,T | N | Y | [`BatchedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedPrepareAndFinalize] |
|
||||
|
||||
|
||||
@@ -159,10 +159,12 @@ Alternatively, you can use the LoRAResolver plugin to dynamically load LoRA adap
|
||||
|
||||
You can set up multiple LoRAResolver plugins if you want to load LoRA adapters from different sources. For example, you might have one resolver for local files and another for S3 storage. vLLM will load the first LoRA adapter that it finds.
|
||||
|
||||
You can either install existing plugins or implement your own. By default, vLLM comes with a [resolver plugin to load LoRA adapters from a local directory.](https://github.com/vllm-project/vllm/tree/main/vllm/plugins/lora_resolvers)
|
||||
To enable this resolver, set `VLLM_ALLOW_RUNTIME_LORA_UPDATING` to True, set `VLLM_PLUGINS` to include `lora_filesystem_resolver`, and then set `VLLM_LORA_RESOLVER_CACHE_DIR` to a local directory. When vLLM receives a request using a LoRA adapter `foobar`,
|
||||
it will first look in the local directory for a directory `foobar`, and attempt to load the contents of that directory as a LoRA adapter. If successful, the request will complete as normal and
|
||||
that adapter will then be available for normal use on the server.
|
||||
You can either install existing plugins or implement your own. By default, vLLM comes with a [resolver plugin to load LoRA adapters from a local directory, as well as a resolver plugin to load LoRA adapters from repositories on Hugging Face Hub](https://github.com/vllm-project/vllm/tree/main/vllm/plugins/lora_resolvers)
|
||||
To enable either of these resolvers, you must `set VLLM_ALLOW_RUNTIME_LORA_UPDATING` to True.
|
||||
|
||||
- To leverage a local directory, set `VLLM_PLUGINS` to include `lora_filesystem_resolver` and set `VLLM_LORA_RESOLVER_CACHE_DIR` to a local directory. When vLLM receives a request using a LoRA adapter `foobar`,
|
||||
it will first look in the local directory for a directory `foobar`, and attempt to load the contents of that directory as a LoRA adapter. If successful, the request will complete as normal and that adapter will then be available for normal use on the server.
|
||||
- To leverage repositories on Hugging Face Hub, set `VLLM_PLUGINS` to include `lora_hf_hub_resolver` and set `VLLM_LORA_RESOLVER_HF_REPO_LIST` to a comma separated list of repository IDs on Hugging Face Hub. When vLLM receives a request for the LoRA adapter `my/repo/subpath`, it will download the adapter at the `subpath` of `my/repo` if it exists and contains an `adapter_config.json`, then build a request to the cached dir for the adapter, similar to the `lora_filesystem_resolver`. Please note that enabling remote downloads is insecure and not intended for use in production environments.
|
||||
|
||||
Alternatively, follow these example steps to implement your own plugin:
|
||||
|
||||
|
||||
@@ -184,6 +184,15 @@ Support use case: Prefill with 'HND' and decode with 'NHD' with experimental con
|
||||
--kv-transfer-config '{..., "enable_permute_local_kv":"True"}'
|
||||
```
|
||||
|
||||
### Cross layers blocks
|
||||
|
||||
By default, this feature is disabled. On attention backends that support this feature, each logical block is contiguous in physical memory. This reduces the number of buffers that need to be transferred.
|
||||
To enable this feature:
|
||||
|
||||
```bash
|
||||
--kv-transfer-config '{..., "kv_connector_extra_config": {"enable_cross_layers_blocks": "True"}}'
|
||||
```
|
||||
|
||||
## Example Scripts/Code
|
||||
|
||||
Refer to these example scripts in the vLLM repository:
|
||||
|
||||
@@ -456,7 +456,6 @@ th {
|
||||
| `StableLmForCausalLM` | StableLM | `stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc. | | |
|
||||
| `Starcoder2ForCausalLM` | Starcoder2 | `bigcode/starcoder2-3b`, `bigcode/starcoder2-7b`, `bigcode/starcoder2-15b`, etc. | | ✅︎ |
|
||||
| `Step1ForCausalLM` | Step-Audio | `stepfun-ai/Step-Audio-EditX`, etc. | ✅︎ | ✅︎ |
|
||||
| `Step3p5ForCausalLM` | Step-3.5-flash | `stepfun-ai/step-3.5-flash`, etc. | | ✅︎ |
|
||||
| `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ |
|
||||
| `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ |
|
||||
| `XverseForCausalLM` | XVERSE | `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -675,6 +674,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `GLM4VForCausalLM`<sup>^</sup> | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ |
|
||||
| `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ |
|
||||
| `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + I<sup>E+</sup> + V<sup>E+</sup> | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ |
|
||||
| `GlmOcrForConditionalGeneration` | GLM-OCR | T + I<sup>E+</sup> | `zai-org/GLM-OCR`, etc. | ✅︎ | ✅︎ |
|
||||
| `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ |
|
||||
| `H2OVLChatModel` | H2OVL | T + I<sup>E+</sup> | `h2oai/h2ovl-mississippi-800m`, `h2oai/h2ovl-mississippi-2b`, etc. | | ✅︎ |
|
||||
| `HunYuanVLForConditionalGeneration` | HunyuanOCR | T + I<sup>E+</sup> | `tencent/HunyuanOCR`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
This example shows how to use vLLM for running offline inference
|
||||
with the correct prompt format on Qwen2.5-Omni (thinker only).
|
||||
with the correct prompt format on Qwen3-Omni (thinker only).
|
||||
"""
|
||||
|
||||
from typing import NamedTuple
|
||||
@@ -112,23 +112,51 @@ def get_multi_audios_query() -> QueryResult:
|
||||
)
|
||||
|
||||
|
||||
def get_multi_images_query() -> QueryResult:
|
||||
question = "What are the differences between these two images?"
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
||||
"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
|
||||
"<|vision_start|><|image_pad|><|vision_end|>"
|
||||
f"{question}<|im_end|>\n"
|
||||
f"<|im_start|>assistant\n"
|
||||
)
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"image": [
|
||||
convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
|
||||
convert_image_mode(ImageAsset("stop_sign").pil_image, "RGB"),
|
||||
],
|
||||
},
|
||||
},
|
||||
limit_mm_per_prompt={
|
||||
"image": 2,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
query_map = {
|
||||
"mixed_modalities": get_mixed_modalities_query,
|
||||
"use_audio_in_video": get_use_audio_in_video_query,
|
||||
"multi_audios": get_multi_audios_query,
|
||||
"multi_images": get_multi_images_query,
|
||||
}
|
||||
|
||||
|
||||
def main(args):
|
||||
model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
|
||||
model_name = args.model
|
||||
query_result = query_map[args.query_type]()
|
||||
|
||||
llm = LLM(
|
||||
model=model_name,
|
||||
max_model_len=12800,
|
||||
max_model_len=args.max_model_len,
|
||||
max_num_seqs=5,
|
||||
limit_mm_per_prompt=query_result.limit_mm_per_prompt,
|
||||
seed=args.seed,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
gpu_memory_utilization=args.gpu_memory_utilization,
|
||||
)
|
||||
|
||||
# We set temperature to 0.2 so that outputs can be different
|
||||
@@ -161,6 +189,31 @@ def parse_args():
|
||||
default=0,
|
||||
help="Set the seed when initializing `vllm.LLM`.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
help="Model name or path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tensor-parallel-size",
|
||||
"-tp",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Tensor parallel size for distributed inference.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gpu-memory-utilization",
|
||||
type=float,
|
||||
default=0.9,
|
||||
help="GPU memory utilization (0.0 to 1.0).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-model-len",
|
||||
type=int,
|
||||
default=12800,
|
||||
help="Maximum model context length.",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -566,6 +566,42 @@ def run_glm4_5v_fp8(questions: list[str], modality: str) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
# GLM-OCR
|
||||
def run_glm_ocr(questions: list[str], modality: str) -> ModelRequestData:
|
||||
model_name = "zai-org/GLM-OCR"
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_model_len=4096,
|
||||
max_num_seqs=2,
|
||||
mm_processor_kwargs={
|
||||
"size": {"shortest_edge": 12544, "longest_edge": 47040000},
|
||||
"fps": 1,
|
||||
},
|
||||
limit_mm_per_prompt={modality: 1},
|
||||
enforce_eager=True,
|
||||
)
|
||||
|
||||
if modality == "image":
|
||||
placeholder = "<|begin_of_image|><|image|><|end_of_image|>"
|
||||
elif modality == "video":
|
||||
placeholder = "<|begin_of_video|><|video|><|end_of_video|>"
|
||||
|
||||
prompts = [
|
||||
(
|
||||
"[gMASK]<sop><|system|>\nYou are a helpful assistant.<|user|>\n"
|
||||
f"{placeholder}"
|
||||
f"{question}<|assistant|>assistant\n"
|
||||
)
|
||||
for question in questions
|
||||
]
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompts=prompts,
|
||||
)
|
||||
|
||||
|
||||
# H2OVL-Mississippi
|
||||
def run_h2ovl(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
@@ -1889,6 +1925,32 @@ def run_step3(questions: list[str], modality: str) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
# StepVL10B
|
||||
def run_step_vl(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
|
||||
model_name = "stepfun-ai/Step3-VL-10B"
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_num_batched_tokens=4096,
|
||||
tensor_parallel_size=1,
|
||||
trust_remote_code=True,
|
||||
limit_mm_per_prompt={modality: 1},
|
||||
reasoning_parser="deepseek_r1",
|
||||
)
|
||||
|
||||
prompts = [
|
||||
"<|begin▁of▁sentence|> You are a helpful assistant.<|BOT|>user\n "
|
||||
f"<im_patch>{question} <|EOT|><|BOT|>assistant\n<think>\n"
|
||||
for question in questions
|
||||
]
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompts=prompts,
|
||||
)
|
||||
|
||||
|
||||
# omni-research/Tarsier-7b
|
||||
def run_tarsier(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
@@ -1962,6 +2024,7 @@ model_example_map = {
|
||||
"glm4_1v": run_glm4_1v,
|
||||
"glm4_5v": run_glm4_5v,
|
||||
"glm4_5v_fp8": run_glm4_5v_fp8,
|
||||
"glm_ocr": run_glm_ocr,
|
||||
"h2ovl_chat": run_h2ovl,
|
||||
"hunyuan_vl": run_hunyuan_vl,
|
||||
"hyperclovax_seed_vision": run_hyperclovax_seed_vision,
|
||||
@@ -2006,6 +2069,7 @@ model_example_map = {
|
||||
"skywork_chat": run_skyworkr1v,
|
||||
"smolvlm": run_smolvlm,
|
||||
"step3": run_step3,
|
||||
"stepvl": run_step_vl,
|
||||
"tarsier": run_tarsier,
|
||||
"tarsier2": run_tarsier2,
|
||||
}
|
||||
@@ -2013,6 +2077,7 @@ model_example_map = {
|
||||
|
||||
MODELS_NEED_VIDEO_METADATA = [
|
||||
"glm4_1v",
|
||||
"glm_ocr",
|
||||
"glm4_5v",
|
||||
"glm4_5v_fp8",
|
||||
"molmo2",
|
||||
|
||||
@@ -1182,6 +1182,32 @@ def load_step3(question: str, image_urls: list[str]) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
def load_step_vl(question: str, image_urls: list[str]) -> ModelRequestData:
|
||||
model_name = "stepfun-ai/Step3-VL-10B"
|
||||
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_num_batched_tokens=4096,
|
||||
limit_mm_per_prompt={"image": len(image_urls)},
|
||||
hf_overrides={"vision_config": {"enable_patch": False}},
|
||||
trust_remote_code=True,
|
||||
reasoning_parser="deepseek_r1",
|
||||
)
|
||||
|
||||
prompt = (
|
||||
"<|begin▁of▁sentence|> You are a helpful assistant.<|BOT|>user\n "
|
||||
f"{'<im_patch>' * len(image_urls)}{question}<|EOT|><|BOT|>"
|
||||
"assistant\n<think>\n"
|
||||
)
|
||||
image_data = [fetch_image(url) for url in image_urls]
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompt=prompt,
|
||||
image_data=image_data,
|
||||
)
|
||||
|
||||
|
||||
def load_tarsier(question: str, image_urls: list[str]) -> ModelRequestData:
|
||||
model_name = "omni-research/Tarsier-7b"
|
||||
|
||||
@@ -1374,6 +1400,7 @@ model_example_map = {
|
||||
"rvl": load_r_vl,
|
||||
"smolvlm": load_smolvlm,
|
||||
"step3": load_step3,
|
||||
"stepvl": load_step_vl,
|
||||
"tarsier": load_tarsier,
|
||||
"tarsier2": load_tarsier2,
|
||||
"glm4_5v": load_glm4_5v,
|
||||
|
||||
@@ -157,6 +157,37 @@ VLLM_CONFIGURE_LOGGING=0 \
|
||||
vllm serve mistralai/Mistral-7B-v0.1 --max-model-len 2048
|
||||
```
|
||||
|
||||
### Example 4: Disable access logs for health check endpoints
|
||||
|
||||
In production environments, health check endpoints like `/health`, `/metrics`,
|
||||
and `/ping` are frequently called by load balancers and monitoring systems,
|
||||
generating a large volume of repetitive access logs. To reduce log noise while
|
||||
keeping logs for other endpoints, use the `--disable-access-log-for-endpoints`
|
||||
option.
|
||||
|
||||
**Disable access logs for health and metrics endpoints:**
|
||||
|
||||
```bash
|
||||
vllm serve mistralai/Mistral-7B-v0.1 --max-model-len 2048 \
|
||||
--disable-access-log-for-endpoints /health,/metrics,/ping
|
||||
```
|
||||
|
||||
**Common endpoints to consider filtering:**
|
||||
|
||||
| Endpoint | Description | Typical Caller |
|
||||
| ---------- | ---------------------- | ---------------------------------------------------- |
|
||||
| `/health` | Health check | Kubernetes liveness/readiness probes, load balancers |
|
||||
| `/metrics` | Prometheus metrics | Prometheus scraper (every 15-60s) |
|
||||
| `/ping` | SageMaker health check | SageMaker infrastructure |
|
||||
| `/load` | Server load metrics | Custom monitoring |
|
||||
|
||||
**Notes:**
|
||||
|
||||
- This option only affects uvicorn access logs, not vLLM application logs
|
||||
- Specify multiple endpoints by separating them with commas (no spaces)
|
||||
- The filter uses exact path matching, query parameters are ignored (e.g., `/health?verbose=true` matches `/health`)
|
||||
- If you need to completely disable all access logs, use `--disable-uvicorn-access-log` instead
|
||||
|
||||
## Additional resources
|
||||
|
||||
- [`logging.config` Dictionary Schema Details](https://docs.python.org/3/library/logging.config.html#dictionary-schema-details)
|
||||
|
||||
@@ -18,32 +18,48 @@ e.g.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import pprint
|
||||
import base64
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
def encode_base64_content_from_url(content_url: str) -> dict[str, str]:
|
||||
"""Encode a content retrieved from a remote url to base64 format."""
|
||||
|
||||
with requests.get(content_url, headers=headers) as response:
|
||||
response.raise_for_status()
|
||||
result = base64.b64encode(response.content).decode("utf-8")
|
||||
|
||||
return {"url": f"data:image/jpeg;base64,{result}"}
|
||||
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
query = "A woman playing with her dog on a beach at sunset."
|
||||
document = (
|
||||
"A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, "
|
||||
"as the dog offers its paw in a heartwarming display of companionship and trust."
|
||||
)
|
||||
image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
documents = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": document,
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(fetch_image(image_url))},
|
||||
},
|
||||
]
|
||||
documents = {
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, "
|
||||
"as the dog offers its paw in a heartwarming display of companionship and trust."
|
||||
),
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": encode_base64_content_from_url(
|
||||
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
),
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
@@ -58,36 +74,23 @@ def main(args):
|
||||
models_url = base_url + "/v1/models"
|
||||
rerank_url = base_url + "/rerank"
|
||||
|
||||
response = requests.get(models_url)
|
||||
response = requests.get(models_url, headers=headers)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
print("Query: string & Document: list of string")
|
||||
prompt = {"model": model, "query": query, "documents": [document]}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: text")
|
||||
prompt = {"model": model, "query": query, "documents": {"content": [documents[0]]}}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: image url")
|
||||
prompt = {
|
||||
data = {
|
||||
"model": model,
|
||||
"query": query,
|
||||
"documents": {"content": [documents[1]]},
|
||||
"documents": documents,
|
||||
}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
response = requests.post(rerank_url, headers=headers, json=data)
|
||||
|
||||
print("Query: string & Document: image base64")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"query": query,
|
||||
"documents": {"content": [documents[2]]},
|
||||
}
|
||||
response = requests.post(rerank_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
# Check the response
|
||||
if response.status_code == 200:
|
||||
print("Request successful!")
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
else:
|
||||
print(f"Request failed with status code: {response.status_code}")
|
||||
print(response.text)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -17,32 +17,48 @@ e.g.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
query = "A woman playing with her dog on a beach at sunset."
|
||||
document = (
|
||||
"A woman shares a joyful moment with her golden retriever on a sun-drenched beach at sunset, "
|
||||
"as the dog offers its paw in a heartwarming display of companionship and trust."
|
||||
)
|
||||
image_url = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"
|
||||
documents = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": document,
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(fetch_image(image_url))},
|
||||
},
|
||||
]
|
||||
def encode_base64_content_from_url(content_url: str) -> dict[str, str]:
|
||||
"""Encode a content retrieved from a remote url to base64 format."""
|
||||
|
||||
with requests.get(content_url, headers=headers) as response:
|
||||
response.raise_for_status()
|
||||
result = base64.b64encode(response.content).decode("utf-8")
|
||||
|
||||
return {"url": f"data:image/jpeg;base64,{result}"}
|
||||
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
queries = "slm markdown"
|
||||
documents = {
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/handelsblatt-preview.png"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": encode_base64_content_from_url(
|
||||
"https://raw.githubusercontent.com/jina-ai/multimodal-reranker-test/main/paper-11.png"
|
||||
),
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
@@ -57,40 +73,15 @@ def main(args):
|
||||
models_url = base_url + "/v1/models"
|
||||
score_url = base_url + "/score"
|
||||
|
||||
response = requests.get(models_url)
|
||||
response = requests.get(models_url, headers=headers)
|
||||
model = response.json()["data"][0]["id"]
|
||||
|
||||
print("Query: string & Document: string")
|
||||
prompt = {"model": model, "queries": query, "documents": document}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: text")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[0]]},
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: image url")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[1]]},
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
|
||||
print("Query: string & Document: image base64")
|
||||
prompt = {
|
||||
"model": model,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[2]]},
|
||||
}
|
||||
response = requests.post(score_url, json=prompt)
|
||||
pprint.pprint(response.json())
|
||||
prompt = {"model": model, "queries": queries, "documents": documents}
|
||||
response = requests.post(score_url, headers=headers, json=prompt)
|
||||
print("\nPrompt when queries is string and documents is a image list:")
|
||||
pprint.pprint(prompt)
|
||||
print("\nScore Response:")
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
+2
-1
@@ -9,7 +9,7 @@ requires = [
|
||||
"torch == 2.9.1",
|
||||
"wheel",
|
||||
"jinja2",
|
||||
"grpcio-tools",
|
||||
"grpcio-tools>=1.76.0",
|
||||
]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
@@ -44,6 +44,7 @@ vllm = "vllm.entrypoints.cli.main:main"
|
||||
|
||||
[project.entry-points."vllm.general_plugins"]
|
||||
lora_filesystem_resolver = "vllm.plugins.lora_resolvers.filesystem_resolver:register_filesystem_resolver"
|
||||
lora_hf_hub_resolver = "vllm.plugins.lora_resolvers.hf_hub_resolver:register_hf_hub_resolver"
|
||||
|
||||
[tool.setuptools_scm]
|
||||
# no extra settings needed, presence enables setuptools-scm
|
||||
|
||||
@@ -9,5 +9,5 @@ wheel
|
||||
jinja2>=3.1.6
|
||||
regex
|
||||
build
|
||||
protobuf >= 6.33.5
|
||||
grpcio-tools
|
||||
protobuf>=6.33.2
|
||||
grpcio-tools>=1.76.0
|
||||
|
||||
@@ -9,9 +9,9 @@ blake3
|
||||
py-cpuinfo
|
||||
transformers >= 4.56.0, < 5
|
||||
tokenizers >= 0.21.1 # Required for fast incremental detokenization.
|
||||
protobuf >= 6.33.5 # Required by LlamaTokenizer, gRPC. CVE-2026-0994
|
||||
protobuf >= 6.30.0 # Required by LlamaTokenizer, gRPC.
|
||||
fastapi[standard] >= 0.115.0 # Required by FastAPI's form models in the OpenAI API server's audio transcriptions endpoint.
|
||||
aiohttp >= 3.13.3
|
||||
aiohttp
|
||||
openai >= 1.99.1 # For Responses API with reasoning content
|
||||
pydantic >= 2.12.0
|
||||
prometheus_client >= 0.18.0
|
||||
@@ -51,5 +51,5 @@ openai-harmony >= 0.0.3 # Required for gpt-oss
|
||||
anthropic >= 0.71.0
|
||||
model-hosting-container-standards >= 0.1.13, < 1.0.0
|
||||
mcp
|
||||
grpcio
|
||||
grpcio-reflection
|
||||
grpcio>=1.76.0
|
||||
grpcio-reflection>=1.76.0
|
||||
@@ -1,2 +1,2 @@
|
||||
lmcache >= 0.3.9
|
||||
lmcache
|
||||
nixl >= 0.7.1 # Required for disaggregated prefill
|
||||
|
||||
@@ -14,7 +14,7 @@ pytest-shard==0.1.2
|
||||
# Async/HTTP dependencies
|
||||
anyio==4.6.2.post1
|
||||
# via httpx, starlette
|
||||
aiohttp==3.13.3
|
||||
aiohttp==3.13.0
|
||||
# via gpt-oss
|
||||
httpx==0.27.2
|
||||
# HTTP testing
|
||||
|
||||
@@ -12,7 +12,7 @@ affine==2.4.0
|
||||
# via rasterio
|
||||
aiohappyeyeballs==2.6.1
|
||||
# via aiohttp
|
||||
aiohttp==3.13.3
|
||||
aiohttp==3.13.0
|
||||
# via
|
||||
# aiohttp-cors
|
||||
# datasets
|
||||
|
||||
@@ -1,48 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from torch._dynamo.utils import counters
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.config import CompilationConfig, CompilationMode, CUDAGraphMode
|
||||
|
||||
|
||||
def test_moe_compilation_cold_start(monkeypatch, use_fresh_inductor_cache):
|
||||
# Run in same process so we can access PyTorch's internal counters
|
||||
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
|
||||
# I'm not sure if this is going to affect the numbers
|
||||
monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0")
|
||||
|
||||
# Force cold compilation
|
||||
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
|
||||
|
||||
compilation_config = CompilationConfig(
|
||||
mode=CompilationMode.VLLM_COMPILE,
|
||||
cudagraph_mode=CUDAGraphMode.NONE, # make the model loading faster
|
||||
)
|
||||
|
||||
counters.clear()
|
||||
|
||||
_ = LLM(
|
||||
model="microsoft/Phi-tiny-MoE-instruct",
|
||||
max_model_len=256,
|
||||
load_format="dummy", # make the model loading faster
|
||||
compilation_config=compilation_config,
|
||||
num_gpu_blocks_override=8, # make the model loading faster
|
||||
)
|
||||
|
||||
# vLLM-compile cold start is special. By default, we do
|
||||
# one full dynamo capture of the entire forward pass.
|
||||
# The forward pass consists of 32 transformer layers.
|
||||
# Then, we split on the attention operation. This results in
|
||||
# 33 subgraphs (not including the attention operation).
|
||||
# The 33 subgraphs then get standalone_compile'd.
|
||||
#
|
||||
# There are actually only 3 unique subgraphs for this model
|
||||
# (all of its transformer layers are the same modulo weights);
|
||||
# this is true for most vLLM models.
|
||||
# So we test that during cold start, the aot_autograd cache
|
||||
# misses for 3 subgraphs and hits for the rest.
|
||||
assert counters["aot_autograd"]["autograd_cache_miss"] == 3
|
||||
assert counters["aot_autograd"]["autograd_cache_hit"] == 30
|
||||
@@ -8,10 +8,6 @@ import torch
|
||||
from torch.fx.experimental.proxy_tensor import make_fx
|
||||
|
||||
from vllm.compilation.backends import split_graph
|
||||
from vllm.compilation.fx_utils import find_op_nodes
|
||||
|
||||
# This import automatically registers `torch.ops.silly.attention`
|
||||
from . import silly_attention # noqa: F401
|
||||
|
||||
|
||||
def test_getitem_moved_to_producer_subgraph():
|
||||
@@ -126,61 +122,3 @@ def test_no_tuple_inputs_with_multiple_consumers():
|
||||
output_split = split_gm(new_x)
|
||||
|
||||
assert torch.allclose(output_original, output_split), "Output mismatch after split"
|
||||
|
||||
|
||||
def test_consecutive_ops_in_split():
|
||||
"""
|
||||
Test that consecutive splitting operations are grouped into the same subgraph
|
||||
"""
|
||||
|
||||
def model_fn(x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Define a simple model where consecutive operations create opportunities
|
||||
for splitting subgraphs.
|
||||
"""
|
||||
# Apply silly attention followed by consecutive operations
|
||||
intermediate = torch.relu(x)
|
||||
attn_inout = torch.sqrt(intermediate)
|
||||
torch.ops.silly.attention(intermediate, intermediate, attn_inout, attn_inout)
|
||||
final_result = torch.sigmoid(attn_inout)
|
||||
return final_result
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
# Create the traced FX graph for the model
|
||||
x = torch.randn(8, 4)
|
||||
|
||||
gm = make_fx(model_fn)(x)
|
||||
|
||||
# Assert presence of the expected operations in the setup
|
||||
assert (
|
||||
len(list(find_op_nodes(torch.ops.aten.relu, gm.graph))) == 1
|
||||
and len(list(find_op_nodes(torch.ops.aten.sqrt, gm.graph))) == 1
|
||||
), "Test setup failed: Expected sqrt and relu operations in the graph."
|
||||
|
||||
# Configure split operations to test
|
||||
splitting_ops = ["silly::attention", "aten::sqrt"]
|
||||
split_gm, split_items = split_graph(gm, splitting_ops)
|
||||
|
||||
# Validate the number of partitions
|
||||
assert len(split_items) == 3, (
|
||||
"Consecutive splitting operations were not grouped correctly."
|
||||
)
|
||||
|
||||
# Validate that correctness is preserved
|
||||
new_x = torch.randn(8, 4)
|
||||
output_original = gm(new_x)
|
||||
output_split = split_gm(new_x)
|
||||
assert torch.allclose(output_original, output_split), (
|
||||
"Output mismatch after splitting."
|
||||
)
|
||||
|
||||
# Check the splitting item has 2 nodes exactly (relu and attn)
|
||||
splitting_items = list(s for s in split_items if s.is_splitting_graph)
|
||||
assert len(splitting_items) == 1, "Expecting a single splitting graph"
|
||||
print(splitting_items[0].graph.graph)
|
||||
splitting_gm = splitting_items[0].graph
|
||||
assert len(splitting_gm.graph.nodes) == 4, "Expecting 4 nodes in splitting graph"
|
||||
assert [node.op for node in splitting_gm.graph.nodes] == ["placeholder"] + 2 * [
|
||||
"call_function"
|
||||
] + ["output"]
|
||||
|
||||
@@ -42,6 +42,7 @@ class MockModelConfig:
|
||||
tokenizer_revision = None
|
||||
multimodal_config = MultiModalConfig()
|
||||
hf_config = MockHFConfig()
|
||||
hf_text_config = MockHFConfig()
|
||||
logits_processor_pattern = None
|
||||
logits_processors: list[str] | None = None
|
||||
diff_sampling_param: dict | None = None
|
||||
|
||||
@@ -518,6 +518,7 @@ class MockModelConfig:
|
||||
tokenizer_revision = None
|
||||
multimodal_config = MultiModalConfig()
|
||||
hf_config = MockHFConfig()
|
||||
hf_text_config = MockHFConfig()
|
||||
logits_processors: list[str] | None = None
|
||||
logits_processor_pattern = None
|
||||
diff_sampling_param: dict | None = None
|
||||
|
||||
@@ -5,9 +5,9 @@ import json
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.entrypoints.test_utils import encode_base64_content_from_url
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.classify.protocol import ClassificationResponse
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
MODEL_NAME = "muziyongshixin/Qwen2.5-VL-7B-for-VideoCls"
|
||||
MAXIMUM_VIDEOS = 1
|
||||
@@ -19,7 +19,7 @@ HF_OVERRIDES = {
|
||||
}
|
||||
input_text = "This product was excellent and exceeded my expectations"
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
image_base64 = {"url": encode_image_url(fetch_image(image_url))}
|
||||
image_base64 = encode_base64_content_from_url(image_url)
|
||||
video_url = "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4"
|
||||
|
||||
|
||||
|
||||
@@ -1,122 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import VLLM_PATH, RemoteOpenAIServer
|
||||
from vllm.entrypoints.pooling.score.protocol import ScoreResponse
|
||||
from vllm.multimodal.utils import encode_image_url, fetch_image
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen3-VL-Reranker-2B"
|
||||
HF_OVERRIDES = {
|
||||
"architectures": ["Qwen3VLForSequenceClassification"],
|
||||
"classifier_from_token": ["no", "yes"],
|
||||
"is_original_qwen3_reranker": True,
|
||||
}
|
||||
|
||||
query = "A cat standing in the snow."
|
||||
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
|
||||
documents = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": query,
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": encode_image_url(fetch_image(image_url))},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"8192",
|
||||
"--chat-template",
|
||||
str(VLLM_PATH / "examples/pooling/score/template/qwen3_vl_reranker.jinja"),
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME, args, override_hf_configs=HF_OVERRIDES
|
||||
) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_str(server: RemoteOpenAIServer):
|
||||
queries = "What is the capital of France?"
|
||||
documents = "The capital of France is Paris."
|
||||
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": queries,
|
||||
"documents": documents,
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_text_content(server: RemoteOpenAIServer):
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[0]]},
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_image_url_content(server: RemoteOpenAIServer):
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[1]]},
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
|
||||
|
||||
def test_score_api_queries_str_documents_image_base64_content(
|
||||
server: RemoteOpenAIServer,
|
||||
):
|
||||
score_response = requests.post(
|
||||
server.url_for("score"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"queries": query,
|
||||
"documents": {"content": [documents[2]]},
|
||||
},
|
||||
)
|
||||
score_response.raise_for_status()
|
||||
score = ScoreResponse.model_validate(score_response.json())
|
||||
|
||||
assert score.id is not None
|
||||
assert score.data is not None
|
||||
assert len(score.data) == 1
|
||||
@@ -1,5 +1,9 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import base64
|
||||
|
||||
import requests
|
||||
|
||||
from vllm.entrypoints.utils import sanitize_message
|
||||
|
||||
|
||||
@@ -8,3 +12,11 @@ def test_sanitize_message():
|
||||
sanitize_message("<_io.BytesIO object at 0x7a95e299e750>")
|
||||
== "<_io.BytesIO object>"
|
||||
)
|
||||
|
||||
|
||||
def encode_base64_content_from_url(content_url: str) -> dict[str, str]:
|
||||
with requests.get(content_url) as response:
|
||||
response.raise_for_status()
|
||||
result = base64.b64encode(response.content).decode("utf-8")
|
||||
|
||||
return {"url": f"data:image/jpeg;base64,{result}"}
|
||||
|
||||
@@ -17,8 +17,6 @@ from vllm.model_executor.layers.activation import (
|
||||
QuickGELU,
|
||||
SiluAndMul,
|
||||
SwigluOAIAndMul,
|
||||
SwigluStepAndMul,
|
||||
swiglustep_and_mul_triton,
|
||||
)
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
@@ -38,7 +36,6 @@ CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 e
|
||||
"gelu_tanh",
|
||||
"fatrelu",
|
||||
"swigluoai_and_mul",
|
||||
"swiglustep_and_mul",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@@ -78,12 +75,9 @@ def test_act_and_mul(
|
||||
elif activation == "swigluoai_and_mul":
|
||||
layer = SwigluOAIAndMul()
|
||||
fn = torch.ops._C.swigluoai_and_mul
|
||||
elif activation == "swiglustep_and_mul":
|
||||
layer = SwigluStepAndMul()
|
||||
fn = swiglustep_and_mul_triton
|
||||
out = layer(x)
|
||||
ref_out = layer.forward_native(x)
|
||||
if activation in ["swigluoai_and_mul", "swiglustep_and_mul"]:
|
||||
if activation == "swigluoai_and_mul":
|
||||
rtol = {
|
||||
# For fp16, change the relative tolerance from 1e-3 to 2e-3
|
||||
torch.float16: 2e-3,
|
||||
@@ -110,7 +104,7 @@ def test_act_and_mul(
|
||||
opcheck(fn, (out, x, threshold))
|
||||
elif activation == "swigluoai_and_mul":
|
||||
opcheck(fn, (out, x, layer.alpha, layer.limit))
|
||||
elif activation != "swiglustep_and_mul":
|
||||
else:
|
||||
opcheck(fn, (out, x))
|
||||
|
||||
|
||||
|
||||
@@ -22,6 +22,9 @@ from vllm.distributed import (
|
||||
)
|
||||
from vllm.forward_context import set_forward_context
|
||||
from vllm.model_executor.layers.fused_moe import fused_topk
|
||||
from vllm.model_executor.layers.fused_moe.all2all_utils import (
|
||||
maybe_make_prepare_finalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEParallelConfig,
|
||||
@@ -40,7 +43,6 @@ from .mk_objects import (
|
||||
TestMoEQuantConfig,
|
||||
expert_info,
|
||||
make_fused_experts,
|
||||
make_prepare_finalize,
|
||||
prepare_finalize_info,
|
||||
)
|
||||
from .parallel_utils import ProcessGroupInfo
|
||||
@@ -603,10 +605,12 @@ def make_modular_kernel(
|
||||
routing_method=RoutingMethodType.DeepSeekV3,
|
||||
)
|
||||
|
||||
# make modular kernel
|
||||
prepare_finalize = make_prepare_finalize(
|
||||
config.prepare_finalize_type, config.all2all_backend(), moe, quant_config
|
||||
prepare_finalize = maybe_make_prepare_finalize(
|
||||
moe=moe,
|
||||
quant_config=quant_config,
|
||||
allow_new_interface=True,
|
||||
)
|
||||
assert prepare_finalize is not None
|
||||
|
||||
fused_experts = make_fused_experts(
|
||||
config.fused_experts_type,
|
||||
|
||||
@@ -7,9 +7,6 @@ import torch
|
||||
# Fused experts and PrepareFinalize imports
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe import TritonExperts
|
||||
from vllm.model_executor.layers.fused_moe.all2all_utils import (
|
||||
maybe_make_prepare_finalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
|
||||
BatchedDeepGemmExperts,
|
||||
)
|
||||
@@ -255,13 +252,12 @@ if has_pplx():
|
||||
)
|
||||
|
||||
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import ( # noqa: E501
|
||||
FlashInferCutlassMoEPrepareAndFinalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
|
||||
FlashInferExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
|
||||
FlashInferCutlassMoEPrepareAndFinalize,
|
||||
create_flashinfer_prepare_finalize,
|
||||
)
|
||||
|
||||
register_prepare_and_finalize(
|
||||
FlashInferCutlassMoEPrepareAndFinalize,
|
||||
@@ -429,24 +425,6 @@ if cutlass_fp4_supported() or has_flashinfer_cutlass_fused_moe():
|
||||
]
|
||||
|
||||
|
||||
def make_prepare_finalize(
|
||||
prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
|
||||
backend: str | None,
|
||||
moe: FusedMoEConfig,
|
||||
quant_config: FusedMoEQuantConfig,
|
||||
) -> mk.FusedMoEPrepareAndFinalize:
|
||||
if backend != "naive" and backend is not None:
|
||||
prepare_finalize = maybe_make_prepare_finalize(moe, quant_config)
|
||||
assert prepare_finalize is not None
|
||||
return prepare_finalize
|
||||
elif prepare_finalize_type == FlashInferCutlassMoEPrepareAndFinalize:
|
||||
return create_flashinfer_prepare_finalize(
|
||||
use_dp=moe.moe_parallel_config.dp_size > 1
|
||||
)
|
||||
else:
|
||||
return MoEPrepareAndFinalizeNoEP()
|
||||
|
||||
|
||||
def _slice(rank: int, num_local_experts: int, t: torch.Tensor) -> torch.Tensor:
|
||||
s = rank * num_local_experts
|
||||
e = s + num_local_experts
|
||||
|
||||
@@ -294,12 +294,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
|
||||
)
|
||||
|
||||
kernel = mk.FusedMoEModularKernel(
|
||||
MoEPrepareAndFinalizeNoEP(
|
||||
defer_input_quant=FlashInferExperts.expects_unquantized_inputs(
|
||||
moe_config=moe_config,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
),
|
||||
MoEPrepareAndFinalizeNoEP(),
|
||||
FlashInferExperts(
|
||||
moe_config=moe_config,
|
||||
quant_config=quant_config,
|
||||
|
||||
@@ -106,12 +106,7 @@ def test_flashinfer_fp4_moe_no_graph(
|
||||
)
|
||||
|
||||
flashinfer_experts = FusedMoEModularKernel(
|
||||
MoEPrepareAndFinalizeNoEP(
|
||||
defer_input_quant=FlashInferExperts.expects_unquantized_inputs(
|
||||
moe_config=moe_config,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
),
|
||||
MoEPrepareAndFinalizeNoEP(),
|
||||
FlashInferExperts(moe_config=moe_config, quant_config=quant_config),
|
||||
)
|
||||
|
||||
|
||||
@@ -715,7 +715,7 @@ def test_mixtral_moe(
|
||||
|
||||
# need to override the forward context for unittests, otherwise it assumes
|
||||
# we're running the model forward pass (the model specified in vllm_config)
|
||||
get_forward_context().all_moe_layers = None
|
||||
get_forward_context().remaining_moe_layers = None
|
||||
|
||||
# Run forward passes for both MoE blocks
|
||||
hf_states, _ = hf_moe.forward(hf_inputs)
|
||||
|
||||
@@ -90,7 +90,7 @@ def test_cutlass_fp4_moe_no_graph(
|
||||
)
|
||||
|
||||
kernel = mk.FusedMoEModularKernel(
|
||||
MoEPrepareAndFinalizeNoEP(defer_input_quant=True),
|
||||
MoEPrepareAndFinalizeNoEP(),
|
||||
CutlassExpertsFp4(
|
||||
moe_config=make_dummy_moe_config(),
|
||||
quant_config=quant_config,
|
||||
|
||||
@@ -87,13 +87,6 @@ NKM_FACTORS_WVSPLITK_FP8 = [
|
||||
SEEDS = [0]
|
||||
|
||||
|
||||
def pad_weights_fp8(weight):
|
||||
num_pad = 256 // weight.element_size()
|
||||
import torch.nn.functional as F
|
||||
|
||||
return F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITKRC)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@@ -198,12 +191,11 @@ def test_rocm_wvsplitk_bias2D_kernel(n, k, m, dtype, seed):
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("padded", [False, True])
|
||||
@pytest.mark.skipif(
|
||||
not (current_platform.is_rocm() and current_platform.supports_fp8()),
|
||||
reason="only test for rocm fp8",
|
||||
)
|
||||
def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed, padded):
|
||||
def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed):
|
||||
torch.manual_seed(seed)
|
||||
|
||||
A = torch.rand(n, k, device="cuda") - 0.5
|
||||
@@ -211,8 +203,6 @@ def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed, padded):
|
||||
|
||||
A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
|
||||
B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
|
||||
if padded:
|
||||
B = pad_weights_fp8(B)
|
||||
|
||||
ref_out = torch._scaled_mm(
|
||||
A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b
|
||||
@@ -232,12 +222,11 @@ def test_rocm_wvsplitk_fp8_kernel(n, k, m, dtype, seed, padded):
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("padded", [False, True])
|
||||
@pytest.mark.skipif(
|
||||
not (current_platform.is_rocm() and current_platform.supports_fp8()),
|
||||
reason="only test for rocm fp8",
|
||||
)
|
||||
def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed, padded):
|
||||
def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed):
|
||||
torch.manual_seed(seed)
|
||||
|
||||
xavier = math.sqrt(2 / k) # normalize to avoid large output-bias deltas
|
||||
@@ -247,8 +236,6 @@ def test_rocm_wvsplitk_fp8_bias1D_kernel(n, k, m, dtype, seed, padded):
|
||||
|
||||
A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
|
||||
B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
|
||||
if padded:
|
||||
B = pad_weights_fp8(B)
|
||||
|
||||
ref_out = torch._scaled_mm(
|
||||
A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b, bias=BIAS
|
||||
|
||||
@@ -458,6 +458,20 @@ VLM_TEST_SETTINGS = {
|
||||
],
|
||||
marks=[large_gpu_mark(min_gb=32)],
|
||||
),
|
||||
"glm_ocr": VLMTestInfo(
|
||||
models=["zai-org/GLM-OCR"],
|
||||
test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE),
|
||||
prompt_formatter=lambda img_prompt: f"[gMASK]<|user|>\n{img_prompt}<|assistant|>\n", # noqa: E501
|
||||
img_idx_to_prompt=lambda idx: "<|begin_of_image|><|image|><|end_of_image|>",
|
||||
video_idx_to_prompt=lambda idx: "<|begin_of_video|><|video|><|end_of_video|>",
|
||||
max_model_len=2048,
|
||||
max_num_seqs=2,
|
||||
get_stop_token_ids=lambda tok: [151329, 151336, 151338],
|
||||
num_logprobs=10,
|
||||
image_size_factors=[(), (0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
|
||||
auto_cls=AutoModelForImageTextToText,
|
||||
marks=[large_gpu_mark(min_gb=32)],
|
||||
),
|
||||
"h2ovl": VLMTestInfo(
|
||||
models=[
|
||||
"h2oai/h2ovl-mississippi-800m",
|
||||
|
||||
@@ -91,6 +91,19 @@ MODEL_CONFIGS: dict[str, dict[str, Any]] = {
|
||||
"use_processor": True,
|
||||
"question": "What is the content of each image?",
|
||||
},
|
||||
"glm_ocr": {
|
||||
"model_name": "zai-org/GLM-OCR",
|
||||
"interface": "llm_generate",
|
||||
"max_model_len": 131072,
|
||||
"max_num_seqs": 2,
|
||||
"sampling_params": {
|
||||
"temperature": 0.0,
|
||||
"max_tokens": 256,
|
||||
"stop_token_ids": None,
|
||||
},
|
||||
"use_processor": True,
|
||||
"question": "Text Recognition:",
|
||||
},
|
||||
"keye_vl": {
|
||||
"model_name": "Kwai-Keye/Keye-VL-8B-Preview",
|
||||
"interface": "llm_generate",
|
||||
|
||||
@@ -122,6 +122,7 @@ MM_DATA_PATCHES = {
|
||||
"ernie4_5_moe_vl": qwen3_vl_patch_mm_data,
|
||||
"glm4v": glm4_1v_patch_mm_data,
|
||||
"glm4v_moe": glm4_1v_patch_mm_data,
|
||||
"glm_ocr": glm4_1v_patch_mm_data,
|
||||
"glmasr": glmasr_patch_mm_data,
|
||||
"molmo2": qwen3_vl_patch_mm_data,
|
||||
"qwen3_vl": qwen3_vl_patch_mm_data,
|
||||
|
||||
+22
-21
@@ -256,7 +256,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
|
||||
),
|
||||
"Exaone4ForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-4.0-32B"),
|
||||
"ExaoneMoEForCausalLM": _HfExamplesInfo(
|
||||
"LGAI-EXAONE/K-EXAONE-236B-A23B", min_transformers_version="5.0.0"
|
||||
"LGAI-EXAONE/K-EXAONE-236B-A23B", min_transformers_version="5.1.0"
|
||||
),
|
||||
"Fairseq2LlamaForCausalLM": _HfExamplesInfo("mgleize/fairseq2-dummy-Llama-3.2-1B"),
|
||||
"FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"),
|
||||
@@ -273,8 +273,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
|
||||
"Glm4MoeForCausalLM": _HfExamplesInfo("zai-org/GLM-4.5"),
|
||||
"Glm4MoeLiteForCausalLM": _HfExamplesInfo(
|
||||
"zai-org/GLM-4.7-Flash",
|
||||
min_transformers_version="5.0.0.dev",
|
||||
is_available_online=False,
|
||||
min_transformers_version="5.0.0",
|
||||
),
|
||||
"GPT2LMHeadModel": _HfExamplesInfo("openai-community/gpt2", {"alias": "gpt2"}),
|
||||
"GPTBigCodeForCausalLM": _HfExamplesInfo(
|
||||
@@ -480,9 +479,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
|
||||
"Step1ForCausalLM": _HfExamplesInfo(
|
||||
"stepfun-ai/Step-Audio-EditX", trust_remote_code=True
|
||||
),
|
||||
"Step3p5ForCausalLM": _HfExamplesInfo(
|
||||
"stepfun-ai/step-3.5-flash", is_available_online=False
|
||||
),
|
||||
"SmolLM3ForCausalLM": _HfExamplesInfo("HuggingFaceTB/SmolLM3-3B"),
|
||||
"StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b"),
|
||||
"StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t"),
|
||||
@@ -654,7 +650,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
# [Decoder-only]
|
||||
"AriaForConditionalGeneration": _HfExamplesInfo("rhymes-ai/Aria"),
|
||||
"AudioFlamingo3ForConditionalGeneration": _HfExamplesInfo(
|
||||
"nvidia/audio-flamingo-3-hf", min_transformers_version="5.0.0.dev"
|
||||
"nvidia/audio-flamingo-3-hf", min_transformers_version="5.0.0"
|
||||
),
|
||||
"AyaVisionForConditionalGeneration": _HfExamplesInfo("CohereLabs/aya-vision-8b"),
|
||||
"BagelForConditionalGeneration": _HfExamplesInfo("ByteDance-Seed/BAGEL-7B-MoT"),
|
||||
@@ -697,7 +693,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
"GlmAsrForConditionalGeneration": _HfExamplesInfo(
|
||||
"zai-org/GLM-ASR-Nano-2512",
|
||||
trust_remote_code=True,
|
||||
min_transformers_version="5.0",
|
||||
min_transformers_version="5.0.0",
|
||||
),
|
||||
"GraniteVision": _HfExamplesInfo("ibm-granite/granite-vision-3.3-2b"),
|
||||
"GraniteSpeechForConditionalGeneration": _HfExamplesInfo(
|
||||
@@ -710,6 +706,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
),
|
||||
"Glm4vForConditionalGeneration": _HfExamplesInfo("zai-org/GLM-4.1V-9B-Thinking"),
|
||||
"Glm4vMoeForConditionalGeneration": _HfExamplesInfo("zai-org/GLM-4.5V"),
|
||||
"GlmOcrForConditionalGeneration": _HfExamplesInfo(
|
||||
"zai-org/GLM-OCR",
|
||||
is_available_online=False,
|
||||
min_transformers_version="5.1.0",
|
||||
),
|
||||
"H2OVLChatModel": _HfExamplesInfo(
|
||||
"h2oai/h2ovl-mississippi-800m",
|
||||
trust_remote_code=True,
|
||||
@@ -1052,7 +1053,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
|
||||
"ExaoneMoeMTP": _HfExamplesInfo(
|
||||
"LGAI-EXAONE/K-EXAONE-236B-A23B",
|
||||
speculative_model="LGAI-EXAONE/K-EXAONE-236B-A23B",
|
||||
min_transformers_version="5.0.0",
|
||||
min_transformers_version="5.1.0",
|
||||
),
|
||||
"Glm4MoeMTPModel": _HfExamplesInfo(
|
||||
"zai-org/GLM-4.5",
|
||||
@@ -1061,7 +1062,13 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
|
||||
"Glm4MoeLiteMTPModel": _HfExamplesInfo(
|
||||
"zai-org/GLM-4.7-Flash",
|
||||
speculative_model="zai-org/GLM-4.7-Flash",
|
||||
min_transformers_version="5.0.0",
|
||||
),
|
||||
"GlmOcrMTPModel": _HfExamplesInfo(
|
||||
"zai-org/GLM-OCR",
|
||||
speculative_model="zai-org/GLM-OCR",
|
||||
is_available_online=False,
|
||||
min_transformers_version="5.1.0",
|
||||
),
|
||||
"LongCatFlashMTPModel": _HfExamplesInfo(
|
||||
"meituan-longcat/LongCat-Flash-Chat",
|
||||
@@ -1084,37 +1091,31 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
|
||||
"Qwen3NextMTP": _HfExamplesInfo(
|
||||
"Qwen/Qwen3-Next-80B-A3B-Instruct", min_transformers_version="4.56.3"
|
||||
),
|
||||
"Step3p5MTP": _HfExamplesInfo(
|
||||
"stepfun-ai/Step-3.5-Flash",
|
||||
trust_remote_code=True,
|
||||
speculative_model="stepfun-ai/Step-3.5-Flash",
|
||||
is_available_online=False,
|
||||
),
|
||||
}
|
||||
|
||||
_TRANSFORMERS_BACKEND_MODELS = {
|
||||
"TransformersEmbeddingModel": _HfExamplesInfo(
|
||||
"BAAI/bge-base-en-v1.5", min_transformers_version="5.0.0.dev"
|
||||
"BAAI/bge-base-en-v1.5", min_transformers_version="5.0.0"
|
||||
),
|
||||
"TransformersForSequenceClassification": _HfExamplesInfo(
|
||||
"papluca/xlm-roberta-base-language-detection",
|
||||
min_transformers_version="5.0.0.dev",
|
||||
min_transformers_version="5.0.0",
|
||||
),
|
||||
"TransformersForCausalLM": _HfExamplesInfo(
|
||||
"hmellor/Ilama-3.2-1B", trust_remote_code=True
|
||||
),
|
||||
"TransformersMultiModalForCausalLM": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),
|
||||
"TransformersMoEForCausalLM": _HfExamplesInfo(
|
||||
"allenai/OLMoE-1B-7B-0924", min_transformers_version="5.0.0.dev"
|
||||
"allenai/OLMoE-1B-7B-0924", min_transformers_version="5.0.0"
|
||||
),
|
||||
"TransformersMultiModalMoEForCausalLM": _HfExamplesInfo(
|
||||
"Qwen/Qwen3-VL-30B-A3B-Instruct", min_transformers_version="5.0.0.dev"
|
||||
"Qwen/Qwen3-VL-30B-A3B-Instruct", min_transformers_version="5.0.0"
|
||||
),
|
||||
"TransformersMoEEmbeddingModel": _HfExamplesInfo(
|
||||
"Qwen/Qwen3-30B-A3B", min_transformers_version="5.0.0.dev"
|
||||
"Qwen/Qwen3-30B-A3B", min_transformers_version="5.0.0"
|
||||
),
|
||||
"TransformersMoEForSequenceClassification": _HfExamplesInfo(
|
||||
"Qwen/Qwen3-30B-A3B", min_transformers_version="5.0.0.dev"
|
||||
"Qwen/Qwen3-30B-A3B", min_transformers_version="5.0.0"
|
||||
),
|
||||
"TransformersMultiModalEmbeddingModel": _HfExamplesInfo("google/gemma-3-4b-it"),
|
||||
"TransformersMultiModalForSequenceClassification": _HfExamplesInfo(
|
||||
|
||||
@@ -78,7 +78,7 @@ def test_models(
|
||||
from packaging.version import Version
|
||||
|
||||
installed = Version(transformers.__version__)
|
||||
required = Version("5.0.0.dev")
|
||||
required = Version("5.0.0")
|
||||
if model == "allenai/OLMoE-1B-7B-0924" and installed < required:
|
||||
pytest.skip(
|
||||
"MoE models with the Transformers modeling backend require "
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from huggingface_hub.constants import HF_HUB_CACHE
|
||||
|
||||
from vllm.plugins.lora_resolvers.hf_hub_resolver import HfHubResolver
|
||||
|
||||
LORA_LIB_MODEL_NAME = "ibm-granite/granite-3.3-8b-instruct"
|
||||
# Repo with multiple LoRAs contained in it
|
||||
LORA_LIB = "ibm-granite/granite-3.3-8b-rag-agent-lib"
|
||||
LORA_NAME = "ibm-granite/granite-3.3-8b-rag-agent-lib/answerability_prediction_lora" # noqa: E501
|
||||
NON_LORA_SUBPATH = "ibm-granite/granite-3.3-8b-rag-agent-lib/README.md"
|
||||
LIB_DOWNLOAD_DIR = os.path.join(
|
||||
HF_HUB_CACHE, "models--ibm-granite--granite-3.3-8b-rag-agent-lib"
|
||||
)
|
||||
INVALID_REPO_NAME = "thisrepodoesnotexist"
|
||||
|
||||
# Repo with only one LoRA in the root dir
|
||||
LORA_REPO_MODEL_NAME = "meta-llama/Llama-2-7b-hf"
|
||||
LORA_REPO = "yard1/llama-2-7b-sql-lora-test"
|
||||
REPO_DOWNLOAD_DIR = os.path.join(
|
||||
HF_HUB_CACHE, "models--yard1--llama-2-7b-sql-lora-test"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hf_resolver_with_direct_path():
|
||||
hf_resolver = HfHubResolver([LORA_REPO])
|
||||
assert hf_resolver is not None
|
||||
|
||||
lora_request = await hf_resolver.resolve_lora(LORA_REPO_MODEL_NAME, LORA_REPO)
|
||||
assert lora_request.lora_name == LORA_REPO
|
||||
assert REPO_DOWNLOAD_DIR in lora_request.lora_path
|
||||
assert "adapter_config.json" in os.listdir(lora_request.lora_path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hf_resolver_with_nested_paths():
|
||||
hf_resolver = HfHubResolver([LORA_LIB])
|
||||
assert hf_resolver is not None
|
||||
|
||||
lora_request = await hf_resolver.resolve_lora(LORA_LIB_MODEL_NAME, LORA_NAME)
|
||||
assert lora_request is not None
|
||||
assert lora_request.lora_name == LORA_NAME
|
||||
assert LIB_DOWNLOAD_DIR in lora_request.lora_path
|
||||
assert "adapter_config.json" in os.listdir(lora_request.lora_path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hf_resolver_with_multiple_repos():
|
||||
hf_resolver = HfHubResolver([LORA_LIB, LORA_REPO])
|
||||
assert hf_resolver is not None
|
||||
|
||||
lora_request = await hf_resolver.resolve_lora(LORA_LIB_MODEL_NAME, LORA_NAME)
|
||||
assert lora_request is not None
|
||||
assert lora_request.lora_name == LORA_NAME
|
||||
assert LIB_DOWNLOAD_DIR in lora_request.lora_path
|
||||
assert "adapter_config.json" in os.listdir(lora_request.lora_path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_adapter():
|
||||
hf_resolver = HfHubResolver([LORA_LIB])
|
||||
assert hf_resolver is not None
|
||||
|
||||
missing_lora_request = await hf_resolver.resolve_lora(LORA_LIB_MODEL_NAME, "foobar")
|
||||
assert missing_lora_request is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nonlora_adapter():
|
||||
hf_resolver = HfHubResolver([LORA_LIB])
|
||||
assert hf_resolver is not None
|
||||
|
||||
readme_request = await hf_resolver.resolve_lora(
|
||||
LORA_LIB_MODEL_NAME, NON_LORA_SUBPATH
|
||||
)
|
||||
assert readme_request is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_repo():
|
||||
hf_resolver = HfHubResolver([LORA_LIB])
|
||||
assert hf_resolver is not None
|
||||
|
||||
invalid_repo_req = await hf_resolver.resolve_lora(
|
||||
INVALID_REPO_NAME,
|
||||
f"{INVALID_REPO_NAME}/foo",
|
||||
)
|
||||
assert invalid_repo_req is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trailing_slash():
|
||||
hf_resolver = HfHubResolver([LORA_LIB])
|
||||
assert hf_resolver is not None
|
||||
|
||||
lora_request = await hf_resolver.resolve_lora(
|
||||
LORA_LIB_MODEL_NAME,
|
||||
f"{LORA_NAME}/",
|
||||
)
|
||||
assert lora_request is not None
|
||||
assert lora_request.lora_name == f"{LORA_NAME}/"
|
||||
assert LIB_DOWNLOAD_DIR in lora_request.lora_path
|
||||
assert "adapter_config.json" in os.listdir(lora_request.lora_path)
|
||||
@@ -36,7 +36,7 @@ class MyGemma2Embedding(nn.Module):
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
input_ids: torch.Tensor | None,
|
||||
positions: torch.Tensor,
|
||||
intermediate_tensors: IntermediateTensors | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
|
||||
@@ -0,0 +1,371 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for the UvicornAccessLogFilter class.
|
||||
"""
|
||||
|
||||
import logging
|
||||
|
||||
from vllm.logging_utils.access_log_filter import (
|
||||
UvicornAccessLogFilter,
|
||||
create_uvicorn_log_config,
|
||||
)
|
||||
|
||||
|
||||
class TestUvicornAccessLogFilter:
|
||||
"""Test cases for UvicornAccessLogFilter."""
|
||||
|
||||
def test_filter_allows_all_when_no_excluded_paths(self):
|
||||
"""Filter should allow all logs when no paths are excluded."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=[])
|
||||
|
||||
record = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "GET", "/v1/completions", "1.1", 200),
|
||||
exc_info=None,
|
||||
)
|
||||
|
||||
assert filter.filter(record) is True
|
||||
|
||||
def test_filter_allows_all_when_excluded_paths_is_none(self):
|
||||
"""Filter should allow all logs when excluded_paths is None."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=None)
|
||||
|
||||
record = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "GET", "/health", "1.1", 200),
|
||||
exc_info=None,
|
||||
)
|
||||
|
||||
assert filter.filter(record) is True
|
||||
|
||||
def test_filter_excludes_health_endpoint(self):
|
||||
"""Filter should exclude /health endpoint when configured."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=["/health"])
|
||||
|
||||
record = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "GET", "/health", "1.1", 200),
|
||||
exc_info=None,
|
||||
)
|
||||
|
||||
assert filter.filter(record) is False
|
||||
|
||||
def test_filter_excludes_metrics_endpoint(self):
|
||||
"""Filter should exclude /metrics endpoint when configured."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=["/metrics"])
|
||||
|
||||
record = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "GET", "/metrics", "1.1", 200),
|
||||
exc_info=None,
|
||||
)
|
||||
|
||||
assert filter.filter(record) is False
|
||||
|
||||
def test_filter_allows_non_excluded_endpoints(self):
|
||||
"""Filter should allow endpoints not in the excluded list."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=["/health", "/metrics"])
|
||||
|
||||
record = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "POST", "/v1/completions", "1.1", 200),
|
||||
exc_info=None,
|
||||
)
|
||||
|
||||
assert filter.filter(record) is True
|
||||
|
||||
def test_filter_excludes_multiple_endpoints(self):
|
||||
"""Filter should exclude multiple configured endpoints."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=["/health", "/metrics", "/ping"])
|
||||
|
||||
# Test /health
|
||||
record_health = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "GET", "/health", "1.1", 200),
|
||||
exc_info=None,
|
||||
)
|
||||
assert filter.filter(record_health) is False
|
||||
|
||||
# Test /metrics
|
||||
record_metrics = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "GET", "/metrics", "1.1", 200),
|
||||
exc_info=None,
|
||||
)
|
||||
assert filter.filter(record_metrics) is False
|
||||
|
||||
# Test /ping
|
||||
record_ping = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "GET", "/ping", "1.1", 200),
|
||||
exc_info=None,
|
||||
)
|
||||
assert filter.filter(record_ping) is False
|
||||
|
||||
def test_filter_with_query_parameters(self):
|
||||
"""Filter should exclude endpoints even with query parameters."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=["/health"])
|
||||
|
||||
record = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "GET", "/health?verbose=true", "1.1", 200),
|
||||
exc_info=None,
|
||||
)
|
||||
|
||||
assert filter.filter(record) is False
|
||||
|
||||
def test_filter_different_http_methods(self):
|
||||
"""Filter should exclude endpoints regardless of HTTP method."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=["/ping"])
|
||||
|
||||
# Test GET
|
||||
record_get = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "GET", "/ping", "1.1", 200),
|
||||
exc_info=None,
|
||||
)
|
||||
assert filter.filter(record_get) is False
|
||||
|
||||
# Test POST
|
||||
record_post = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "POST", "/ping", "1.1", 200),
|
||||
exc_info=None,
|
||||
)
|
||||
assert filter.filter(record_post) is False
|
||||
|
||||
def test_filter_with_different_status_codes(self):
|
||||
"""Filter should exclude endpoints regardless of status code."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=["/health"])
|
||||
|
||||
for status_code in [200, 500, 503]:
|
||||
record = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg='%s - "%s %s HTTP/%s" %d',
|
||||
args=("127.0.0.1:12345", "GET", "/health", "1.1", status_code),
|
||||
exc_info=None,
|
||||
)
|
||||
assert filter.filter(record) is False
|
||||
|
||||
|
||||
class TestCreateUvicornLogConfig:
|
||||
"""Test cases for create_uvicorn_log_config function."""
|
||||
|
||||
def test_creates_valid_config_structure(self):
|
||||
"""Config should have required logging configuration keys."""
|
||||
config = create_uvicorn_log_config(excluded_paths=["/health"])
|
||||
|
||||
assert "version" in config
|
||||
assert config["version"] == 1
|
||||
assert "disable_existing_loggers" in config
|
||||
assert "formatters" in config
|
||||
assert "handlers" in config
|
||||
assert "loggers" in config
|
||||
assert "filters" in config
|
||||
|
||||
def test_config_includes_access_log_filter(self):
|
||||
"""Config should include the access log filter."""
|
||||
config = create_uvicorn_log_config(excluded_paths=["/health", "/metrics"])
|
||||
|
||||
assert "access_log_filter" in config["filters"]
|
||||
filter_config = config["filters"]["access_log_filter"]
|
||||
assert filter_config["()"] == UvicornAccessLogFilter
|
||||
assert filter_config["excluded_paths"] == ["/health", "/metrics"]
|
||||
|
||||
def test_config_applies_filter_to_access_handler(self):
|
||||
"""Config should apply the filter to the access handler."""
|
||||
config = create_uvicorn_log_config(excluded_paths=["/health"])
|
||||
|
||||
assert "access" in config["handlers"]
|
||||
assert "filters" in config["handlers"]["access"]
|
||||
assert "access_log_filter" in config["handlers"]["access"]["filters"]
|
||||
|
||||
def test_config_with_custom_log_level(self):
|
||||
"""Config should respect custom log level."""
|
||||
config = create_uvicorn_log_config(
|
||||
excluded_paths=["/health"], log_level="debug"
|
||||
)
|
||||
|
||||
assert config["loggers"]["uvicorn"]["level"] == "DEBUG"
|
||||
assert config["loggers"]["uvicorn.access"]["level"] == "DEBUG"
|
||||
assert config["loggers"]["uvicorn.error"]["level"] == "DEBUG"
|
||||
|
||||
def test_config_with_empty_excluded_paths(self):
|
||||
"""Config should work with empty excluded paths."""
|
||||
config = create_uvicorn_log_config(excluded_paths=[])
|
||||
|
||||
assert config["filters"]["access_log_filter"]["excluded_paths"] == []
|
||||
|
||||
def test_config_with_none_excluded_paths(self):
|
||||
"""Config should work with None excluded paths."""
|
||||
config = create_uvicorn_log_config(excluded_paths=None)
|
||||
|
||||
assert config["filters"]["access_log_filter"]["excluded_paths"] == []
|
||||
|
||||
|
||||
class TestIntegration:
|
||||
"""Integration tests for the access log filter."""
|
||||
|
||||
def test_filter_with_real_logger(self):
|
||||
"""Test filter works with a real Python logger simulating uvicorn."""
|
||||
# Create a logger with our filter (simulating uvicorn.access)
|
||||
logger = logging.getLogger("uvicorn.access")
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
# Clear any existing handlers
|
||||
logger.handlers = []
|
||||
|
||||
# Create a custom handler that tracks messages
|
||||
logged_messages: list[str] = []
|
||||
|
||||
class TrackingHandler(logging.Handler):
|
||||
def emit(self, record):
|
||||
logged_messages.append(record.getMessage())
|
||||
|
||||
handler = TrackingHandler()
|
||||
handler.setLevel(logging.INFO)
|
||||
filter = UvicornAccessLogFilter(excluded_paths=["/health", "/metrics"])
|
||||
handler.addFilter(filter)
|
||||
logger.addHandler(handler)
|
||||
|
||||
# Log using uvicorn's format with args tuple
|
||||
# Format: '%s - "%s %s HTTP/%s" %d'
|
||||
logger.info(
|
||||
'%s - "%s %s HTTP/%s" %d',
|
||||
"127.0.0.1:12345",
|
||||
"GET",
|
||||
"/health",
|
||||
"1.1",
|
||||
200,
|
||||
)
|
||||
logger.info(
|
||||
'%s - "%s %s HTTP/%s" %d',
|
||||
"127.0.0.1:12345",
|
||||
"GET",
|
||||
"/v1/completions",
|
||||
"1.1",
|
||||
200,
|
||||
)
|
||||
logger.info(
|
||||
'%s - "%s %s HTTP/%s" %d',
|
||||
"127.0.0.1:12345",
|
||||
"GET",
|
||||
"/metrics",
|
||||
"1.1",
|
||||
200,
|
||||
)
|
||||
logger.info(
|
||||
'%s - "%s %s HTTP/%s" %d',
|
||||
"127.0.0.1:12345",
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
"1.1",
|
||||
200,
|
||||
)
|
||||
|
||||
# Verify only non-excluded endpoints were logged
|
||||
assert len(logged_messages) == 2
|
||||
assert "/v1/completions" in logged_messages[0]
|
||||
assert "/v1/chat/completions" in logged_messages[1]
|
||||
|
||||
def test_filter_allows_non_uvicorn_access_logs(self):
|
||||
"""Test filter allows logs from non-uvicorn.access loggers."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=["/health"])
|
||||
|
||||
# Log record from a different logger name
|
||||
record = logging.LogRecord(
|
||||
name="uvicorn.error",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg="Some error message about /health",
|
||||
args=(),
|
||||
exc_info=None,
|
||||
)
|
||||
|
||||
# Should allow because it's not from uvicorn.access
|
||||
assert filter.filter(record) is True
|
||||
|
||||
def test_filter_handles_malformed_args(self):
|
||||
"""Test filter handles log records with unexpected args format."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=["/health"])
|
||||
|
||||
# Log record with insufficient args
|
||||
record = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg="Some message",
|
||||
args=("only", "two"),
|
||||
exc_info=None,
|
||||
)
|
||||
|
||||
# Should allow because args doesn't have expected format
|
||||
assert filter.filter(record) is True
|
||||
|
||||
def test_filter_handles_non_tuple_args(self):
|
||||
"""Test filter handles log records with non-tuple args."""
|
||||
filter = UvicornAccessLogFilter(excluded_paths=["/health"])
|
||||
|
||||
# Log record with None args
|
||||
record = logging.LogRecord(
|
||||
name="uvicorn.access",
|
||||
level=logging.INFO,
|
||||
pathname="",
|
||||
lineno=0,
|
||||
msg="Some message without args",
|
||||
args=None,
|
||||
exc_info=None,
|
||||
)
|
||||
|
||||
# Should allow because args is None
|
||||
assert filter.filter(record) is True
|
||||
@@ -107,10 +107,7 @@ def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
|
||||
|
||||
|
||||
def make_kv_cache_config_hybrid_model(
|
||||
block_size: int,
|
||||
num_blocks: int,
|
||||
sliding_window_blocks: int,
|
||||
second_spec_type: str = "sliding_window",
|
||||
block_size: int, num_blocks: int, second_spec_type: str = "sliding_window"
|
||||
) -> KVCacheConfig:
|
||||
if second_spec_type == "sliding_window":
|
||||
second_spec = SlidingWindowSpec(
|
||||
@@ -118,7 +115,7 @@ def make_kv_cache_config_hybrid_model(
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=sliding_window_blocks * block_size,
|
||||
sliding_window=2 * block_size,
|
||||
)
|
||||
elif second_spec_type == "mamba":
|
||||
second_spec = MambaSpec(
|
||||
@@ -328,7 +325,7 @@ def test_prefill(hash_fn):
|
||||
def test_prefill_hybrid_model():
|
||||
block_size = 16
|
||||
manager = KVCacheManager(
|
||||
make_kv_cache_config_hybrid_model(block_size, 21, 2),
|
||||
make_kv_cache_config_hybrid_model(block_size, 21),
|
||||
max_model_len=8192,
|
||||
enable_caching=True,
|
||||
hash_block_size=block_size,
|
||||
@@ -337,8 +334,7 @@ def test_prefill_hybrid_model():
|
||||
hash_fn = sha256
|
||||
|
||||
# Complete 3 blocks (48 tokens)
|
||||
num_full_blocks = 3
|
||||
common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]
|
||||
common_token_ids = [i for i in range(3) for _ in range(block_size)]
|
||||
|
||||
# Fully cache miss
|
||||
# Incomplete 1 block (7 tokens)
|
||||
@@ -379,7 +375,6 @@ def test_prefill_hybrid_model():
|
||||
# Cache hit in the common prefix
|
||||
# Incomplete 1 block (5 tokens)
|
||||
unique_token_ids = [3] * 5
|
||||
all_token_ids = common_token_ids + unique_token_ids
|
||||
req1 = make_request("1", common_token_ids + unique_token_ids, block_size, hash_fn)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
|
||||
assert len(req1.block_hashes) == 3
|
||||
@@ -399,13 +394,34 @@ def test_prefill_hybrid_model():
|
||||
manager.free(req0)
|
||||
manager.free(req1)
|
||||
|
||||
cached_block_hash_to_block_bak = copy.copy(
|
||||
manager.block_pool.cached_block_hash_to_block._cache
|
||||
)
|
||||
|
||||
def test_partial_request_hit(
|
||||
request_id: str,
|
||||
hash_to_evict: list[BlockHashWithGroupId],
|
||||
expect_hit_length: int,
|
||||
):
|
||||
req = make_request(
|
||||
request_id, common_token_ids + unique_token_ids, block_size, sha256
|
||||
)
|
||||
for hash_with_group_id in hash_to_evict:
|
||||
manager.block_pool.cached_block_hash_to_block._cache.pop(hash_with_group_id)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
|
||||
assert len(req.block_hashes) == 3
|
||||
assert num_computed_tokens == expect_hit_length * block_size
|
||||
for block_per_group in computed_blocks.blocks:
|
||||
assert len(block_per_group) == num_computed_tokens // block_size
|
||||
for hash_with_group_id in hash_to_evict:
|
||||
manager.block_pool.cached_block_hash_to_block._cache[hash_with_group_id] = (
|
||||
cached_block_hash_to_block_bak[hash_with_group_id]
|
||||
)
|
||||
manager.free(req)
|
||||
|
||||
# Evict the blocks outside sliding window, does not affect the hit length.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
test_partial_request_hit(
|
||||
"2",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[0], 1),
|
||||
make_block_hash_with_group_id(block_hashes[0], 2),
|
||||
@@ -414,23 +430,13 @@ def test_prefill_hybrid_model():
|
||||
)
|
||||
|
||||
# Evict the first block of full attention, makes total cache miss.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"3",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[0], 0)],
|
||||
0,
|
||||
test_partial_request_hit(
|
||||
"3", [make_block_hash_with_group_id(block_hashes[0], 0)], 0
|
||||
)
|
||||
|
||||
# Evict the last block of all layers, reduces the hit length to 2.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
test_partial_request_hit(
|
||||
"4",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[2], 0),
|
||||
make_block_hash_with_group_id(block_hashes[2], 1),
|
||||
@@ -440,36 +446,18 @@ def test_prefill_hybrid_model():
|
||||
)
|
||||
|
||||
# Evict the last block of full attention, reduces the hit length to 2.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"5",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[2], 0)],
|
||||
2,
|
||||
test_partial_request_hit(
|
||||
"5", [make_block_hash_with_group_id(block_hashes[2], 0)], 2
|
||||
)
|
||||
|
||||
# Evict the last block of sliding window, reduces the hit length to 2.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"6",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[2], 1)],
|
||||
2,
|
||||
test_partial_request_hit(
|
||||
"6", [make_block_hash_with_group_id(block_hashes[2], 1)], 2
|
||||
)
|
||||
|
||||
# Evict the last block of sliding window, reduces the hit length to 2.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"7",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[2], 2)],
|
||||
2,
|
||||
test_partial_request_hit(
|
||||
"7", [make_block_hash_with_group_id(block_hashes[2], 2)], 2
|
||||
)
|
||||
|
||||
# Evict different set of blocks for full attention and sliding window makes
|
||||
@@ -478,12 +466,8 @@ def test_prefill_hybrid_model():
|
||||
# The cache hit length of sliding window is 2 * block_size.
|
||||
# Then it is cache miss as the two type of layers
|
||||
# have different hit length.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
test_partial_request_hit(
|
||||
"8",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[2], 0),
|
||||
make_block_hash_with_group_id(block_hashes[0], 1),
|
||||
@@ -493,214 +477,6 @@ def test_prefill_hybrid_model():
|
||||
)
|
||||
|
||||
|
||||
def test_prefill_hybrid_model_eagle():
|
||||
block_size = 16
|
||||
kv_cache_config = make_kv_cache_config_hybrid_model(block_size, 31, 3)
|
||||
manager = KVCacheManager(
|
||||
kv_cache_config,
|
||||
max_model_len=8192,
|
||||
enable_caching=True,
|
||||
hash_block_size=block_size,
|
||||
use_eagle=True,
|
||||
)
|
||||
|
||||
hash_fn = sha256
|
||||
|
||||
# Complete 6 blocks (96 tokens)
|
||||
num_full_blocks = 6
|
||||
common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]
|
||||
|
||||
# Fully cache miss
|
||||
# Incomplete 1 block (7 tokens)
|
||||
unique_token_ids = [6] * 7
|
||||
all_token_ids = common_token_ids + unique_token_ids
|
||||
req0 = make_request("0", all_token_ids, block_size, hash_fn)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
|
||||
assert len(req0.block_hashes) == len(all_token_ids) // block_size
|
||||
assert not computed_blocks.blocks[0]
|
||||
assert num_computed_tokens == 0
|
||||
blocks = manager.allocate_slots(
|
||||
req0, len(all_token_ids), num_computed_tokens, computed_blocks
|
||||
)
|
||||
block_ids = (
|
||||
[1, 2, 3, 4, 5, 6, 7],
|
||||
[8, 9, 10, 11, 12, 13, 14],
|
||||
[15, 16, 17, 18, 19, 20, 21],
|
||||
)
|
||||
assert blocks is not None and blocks.get_block_ids() == block_ids
|
||||
|
||||
# Check full block metadata
|
||||
parent_block_hash = None
|
||||
for i, full_block_ids in enumerate(zip(*(row[:-1] for row in block_ids))):
|
||||
block_tokens = tuple(all_token_ids[i * block_size : (i + 1) * block_size])
|
||||
block_hash = hash_block_tokens(hash_fn, parent_block_hash, block_tokens)
|
||||
for group_id, block_id in enumerate(full_block_ids):
|
||||
blk_hash = manager.block_pool.blocks[block_id].block_hash
|
||||
assert blk_hash is not None
|
||||
assert get_block_hash(blk_hash) == block_hash
|
||||
assert get_group_id(blk_hash) == group_id
|
||||
assert manager.block_pool.blocks[block_id].ref_cnt == 1
|
||||
parent_block_hash = block_hash
|
||||
|
||||
# Check partial block metadata
|
||||
for partial_block_id in (row[-1] for row in block_ids):
|
||||
assert manager.block_pool.blocks[partial_block_id].block_hash is None
|
||||
assert manager.block_pool.blocks[partial_block_id].ref_cnt == 1
|
||||
|
||||
# Cache hit in the common prefix
|
||||
# Incomplete 1 block (5 tokens)
|
||||
unique_token_ids = [6] * 5
|
||||
all_token_ids = common_token_ids + unique_token_ids
|
||||
req1 = make_request("1", all_token_ids, block_size, hash_fn)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
|
||||
assert len(req1.block_hashes) == num_full_blocks
|
||||
assert computed_blocks.get_block_ids() == (
|
||||
[1, 2, 3, 4],
|
||||
[0, 9, 10, 11],
|
||||
[0, 16, 17, 18],
|
||||
)
|
||||
assert num_computed_tokens == 4 * block_size
|
||||
num_new_tokens = len(all_token_ids) - num_computed_tokens
|
||||
blocks = manager.allocate_slots(
|
||||
req1, num_new_tokens, num_computed_tokens, computed_blocks
|
||||
)
|
||||
assert blocks is not None and blocks.get_block_ids() == (
|
||||
[22, 23, 24],
|
||||
[25, 26, 27],
|
||||
[28, 29, 30],
|
||||
)
|
||||
for block_per_group in computed_blocks.blocks:
|
||||
for block in block_per_group:
|
||||
if block != manager.block_pool.null_block:
|
||||
assert block.ref_cnt == 2
|
||||
|
||||
block_hashes = req1.block_hashes
|
||||
manager.free(req0)
|
||||
manager.free(req1)
|
||||
|
||||
# Evict the blocks outside sliding window, does not affect the hit length.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"2",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[0], 1),
|
||||
make_block_hash_with_group_id(block_hashes[0], 2),
|
||||
],
|
||||
4,
|
||||
)
|
||||
|
||||
# Evict the first block of full attention, makes total cache miss.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"3",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[0], 0)],
|
||||
0,
|
||||
)
|
||||
|
||||
# Evict the last block of all layers, reduces the hit length to 3.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"4",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[-1], 0),
|
||||
make_block_hash_with_group_id(block_hashes[-1], 1),
|
||||
make_block_hash_with_group_id(block_hashes[-1], 2),
|
||||
],
|
||||
3,
|
||||
)
|
||||
|
||||
# Evict the last block of full attention, reduces the hit length to 3.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"5",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[-1], 0)],
|
||||
3,
|
||||
)
|
||||
|
||||
# Since the last block of full attention is dropped for eagle, evict
|
||||
# the second last block of sliding window, reduces the hit length to 3.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"6",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[-2], 1)],
|
||||
3,
|
||||
)
|
||||
|
||||
# Since the last block of full attention is dropped for eagle, evict
|
||||
# the second last block of sliding window, reduces the hit length to 3.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"7",
|
||||
all_token_ids,
|
||||
[make_block_hash_with_group_id(block_hashes[-2], 2)],
|
||||
3,
|
||||
)
|
||||
|
||||
# Evict different set of blocks for full attention and sliding window makes
|
||||
# total cache miss.
|
||||
# The cache hit length of full attention is 4 * block_size.
|
||||
# The cache hit length of sliding window is 3 * block_size.
|
||||
# Then it is cache miss as the two type of layers
|
||||
# have different hit length.
|
||||
_test_partial_request_hit(
|
||||
manager,
|
||||
block_size,
|
||||
num_full_blocks,
|
||||
"8",
|
||||
all_token_ids,
|
||||
[
|
||||
make_block_hash_with_group_id(block_hashes[-1], 0),
|
||||
make_block_hash_with_group_id(block_hashes[0], 1),
|
||||
make_block_hash_with_group_id(block_hashes[0], 2),
|
||||
],
|
||||
0,
|
||||
)
|
||||
|
||||
|
||||
def _test_partial_request_hit(
|
||||
manager: KVCacheManager,
|
||||
block_size: int,
|
||||
num_full_blocks,
|
||||
request_id: str,
|
||||
prompt_token_ids: list[int],
|
||||
hash_to_evict: list[BlockHashWithGroupId],
|
||||
expect_hit_length: int,
|
||||
):
|
||||
cached_block_hash_to_block_bak = copy.copy(
|
||||
manager.block_pool.cached_block_hash_to_block._cache
|
||||
)
|
||||
req = make_request(request_id, prompt_token_ids, block_size, sha256)
|
||||
for hash_with_group_id in hash_to_evict:
|
||||
manager.block_pool.cached_block_hash_to_block._cache.pop(hash_with_group_id)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
|
||||
assert len(req.block_hashes) == num_full_blocks
|
||||
assert num_computed_tokens == expect_hit_length * block_size
|
||||
for block_per_group in computed_blocks.blocks:
|
||||
assert len(block_per_group) == num_computed_tokens // block_size
|
||||
for hash_with_group_id in hash_to_evict:
|
||||
manager.block_pool.cached_block_hash_to_block._cache[hash_with_group_id] = (
|
||||
cached_block_hash_to_block_bak[hash_with_group_id]
|
||||
)
|
||||
manager.free(req)
|
||||
|
||||
|
||||
def _make_hybrid_kv_cache_config(
|
||||
block_size: int, num_blocks: int, spec_types: list[str]
|
||||
) -> KVCacheConfig:
|
||||
@@ -879,85 +655,6 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]):
|
||||
manager.free(req1)
|
||||
|
||||
|
||||
# Test cases with eagle enabled: Only test a single simple case for now.
|
||||
# - 2 groups: 1 full + 1 other
|
||||
_EAGLE_HYBRID_MODEL_TEST_CASES = [
|
||||
# 2 groups: 1 full + 1 other
|
||||
pytest.param(["full", "sliding_window"], 2, id="2g-full+sw"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spec_types,expect_hit_length", _EAGLE_HYBRID_MODEL_TEST_CASES)
|
||||
def test_prefill_hybrid_model_combinations_eagle(
|
||||
spec_types: list[str], expect_hit_length: int
|
||||
):
|
||||
"""
|
||||
Test prefix caching with hybrid models (1 full attn + 1 other) with EAGLE.
|
||||
More complex hybrid models with EAGLE are not yet supported (see issue #32802).
|
||||
"""
|
||||
block_size = 16
|
||||
num_groups = len(spec_types)
|
||||
# Allocate enough blocks for all groups
|
||||
num_blocks = 10 * num_groups
|
||||
|
||||
kv_cache_config = _make_hybrid_kv_cache_config(block_size, num_blocks, spec_types)
|
||||
manager = KVCacheManager(
|
||||
kv_cache_config,
|
||||
max_model_len=8192,
|
||||
enable_caching=True,
|
||||
hash_block_size=block_size,
|
||||
use_eagle=True,
|
||||
)
|
||||
|
||||
hash_fn = sha256
|
||||
|
||||
# Complete 3 blocks (48 tokens)
|
||||
num_full_blocks = 4
|
||||
common_token_ids = [i for i in range(num_full_blocks) for _ in range(block_size)]
|
||||
unique_token_ids = [4] * 7
|
||||
all_token_ids = common_token_ids + unique_token_ids
|
||||
|
||||
# First request: no cache hit initially
|
||||
req0 = make_request("0", all_token_ids, block_size, hash_fn)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
|
||||
|
||||
assert len(req0.block_hashes) == num_full_blocks
|
||||
assert not computed_blocks.blocks[0] # No cache hit initially
|
||||
assert num_computed_tokens == 0
|
||||
|
||||
blocks = manager.allocate_slots(
|
||||
req0, len(all_token_ids), num_computed_tokens, computed_blocks
|
||||
)
|
||||
assert blocks is not None
|
||||
# Should have blocks for all groups
|
||||
assert len(blocks.get_block_ids()) == num_groups
|
||||
|
||||
# Second request: should hit cached blocks for common prefix
|
||||
all_token_ids = common_token_ids + [6] * 5
|
||||
req1 = make_request("1", all_token_ids, block_size, hash_fn)
|
||||
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
|
||||
|
||||
# Should hit cached blocks for all groups
|
||||
assert num_computed_tokens == expect_hit_length * block_size
|
||||
assert len(computed_blocks.blocks) == num_groups
|
||||
# Verify each group has the correct number of computed blocks
|
||||
for block_per_group in computed_blocks.blocks:
|
||||
assert len(block_per_group) == expect_hit_length
|
||||
|
||||
# Allocate and verify blocks for second request
|
||||
blocks = manager.allocate_slots(
|
||||
req1,
|
||||
len(all_token_ids) - num_computed_tokens,
|
||||
num_computed_tokens,
|
||||
computed_blocks,
|
||||
)
|
||||
assert blocks is not None
|
||||
assert len(blocks.get_block_ids()) == num_groups
|
||||
|
||||
manager.free(req0)
|
||||
manager.free(req1)
|
||||
|
||||
|
||||
def test_prefill_plp():
|
||||
"""Test prefill with APC and some prompt logprobs (plp) requests.
|
||||
|
||||
|
||||
@@ -870,66 +870,6 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):
|
||||
assert stats.num_accepted_tokens_per_pos == expected[3]
|
||||
|
||||
|
||||
def test_spec_decoding_stats_empty_output():
|
||||
"""Test that spec decoding stats handle empty output tokens gracefully.
|
||||
|
||||
This is a regression test for a bug where empty sampled_token_ids
|
||||
would cause num_accepted = len([]) - 1 = -1, leading to a
|
||||
ValueError when incrementing a Prometheus counter with a negative value.
|
||||
"""
|
||||
num_spec_tokens = 3
|
||||
scheduler = create_scheduler(num_speculative_tokens=num_spec_tokens)
|
||||
requests = create_requests(num_requests=1, num_tokens=1)
|
||||
request = requests[0]
|
||||
req_id = request.request_id
|
||||
|
||||
scheduler.add_request(request)
|
||||
|
||||
# Initial schedule (prefill)
|
||||
output = scheduler.schedule()
|
||||
assert len(output.scheduled_new_reqs) == 1
|
||||
|
||||
# Complete the prefill with a sampled token
|
||||
model_runner_output = ModelRunnerOutput(
|
||||
req_ids=[req_id],
|
||||
req_id_to_index={req_id: 0},
|
||||
sampled_token_ids=[[0]],
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=[],
|
||||
)
|
||||
scheduler.update_from_output(output, model_runner_output)
|
||||
|
||||
# Add draft tokens for speculation
|
||||
draft_token_ids = DraftTokenIds([req_id], [[1, 2, 3]])
|
||||
scheduler.update_draft_token_ids(draft_token_ids)
|
||||
|
||||
# Schedule the speculated tokens for validation
|
||||
output = scheduler.schedule()
|
||||
assert req_id in output.scheduled_spec_decode_tokens
|
||||
assert len(output.scheduled_spec_decode_tokens[req_id]) == 3
|
||||
|
||||
# Simulate empty output tokens (e.g., due to request abortion or error)
|
||||
# This would previously cause num_accepted = -1 and crash
|
||||
model_runner_output = ModelRunnerOutput(
|
||||
req_ids=[req_id],
|
||||
req_id_to_index={req_id: 0},
|
||||
sampled_token_ids=[[]], # Empty output tokens
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=[],
|
||||
)
|
||||
|
||||
# This should not raise an error
|
||||
engine_core_outputs = scheduler.update_from_output(output, model_runner_output)
|
||||
|
||||
# Spec decoding stats should be None since no tokens were generated
|
||||
scheduler_stats = (
|
||||
engine_core_outputs[0].scheduler_stats if engine_core_outputs else None
|
||||
)
|
||||
assert scheduler_stats is None or scheduler_stats.spec_decoding_stats is None
|
||||
|
||||
|
||||
def _assert_right_scheduler_output(
|
||||
output: SchedulerOutput,
|
||||
num_requests: int,
|
||||
|
||||
@@ -455,7 +455,7 @@ def test_eagle_correctness(
|
||||
from packaging.version import Version
|
||||
|
||||
installed = Version(transformers.__version__)
|
||||
required = Version("5.0.0.dev")
|
||||
required = Version("5.0.0")
|
||||
if installed < required:
|
||||
pytest.skip(
|
||||
"Eagle3 with the Transformers modeling backend requires "
|
||||
|
||||
@@ -34,11 +34,18 @@ else
|
||||
KV_CONFIG_HETERO_LAYOUT=''
|
||||
fi
|
||||
|
||||
CROSS_LAYERS_BLOCKS=${CROSS_LAYERS_BLOCKS:-"False"} # Default to non cross layers
|
||||
if [[ "$CROSS_LAYERS_BLOCKS" == "True" ]]; then
|
||||
KV_EXTRA_CONFIG=',"kv_connector_extra_config":{"cross_layers_blocks": "True"}'
|
||||
else
|
||||
KV_EXTRA_CONFIG=''
|
||||
fi
|
||||
|
||||
# Build the kv-transfer-config once
|
||||
if [[ "$KV_BUFFER_DEVICE" == "cuda" ]]; then
|
||||
KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}'}'
|
||||
KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"'${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}'}'
|
||||
else
|
||||
KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}"}"
|
||||
KV_CONFIG="{\"kv_connector\":\"NixlConnector\",\"kv_role\":\"kv_both\",\"kv_buffer_device\":\"$KV_BUFFER_DEVICE\""${KV_CONFIG_HETERO_LAYOUT}${KV_EXTRA_CONFIG}"}"
|
||||
fi
|
||||
|
||||
# Models to run
|
||||
|
||||
@@ -18,8 +18,12 @@ import ray
|
||||
import torch
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.config import KVTransferConfig
|
||||
from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
|
||||
from vllm.config import KVTransferConfig, set_current_vllm_config
|
||||
from vllm.distributed.kv_transfer.kv_connector.utils import (
|
||||
KVOutputAggregator,
|
||||
TpKVTopology,
|
||||
get_current_attn_backend,
|
||||
)
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1 import nixl_connector
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
|
||||
@@ -48,8 +52,11 @@ from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.output_processor import OutputProcessor
|
||||
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor
|
||||
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
|
||||
from vllm.v1.request import RequestStatus
|
||||
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
|
||||
from vllm.v1.worker.utils import AttentionGroup
|
||||
|
||||
from .utils import create_request, create_scheduler, create_vllm_config
|
||||
|
||||
@@ -366,6 +373,7 @@ def test_kv_transfer_handshake(dist_init):
|
||||
|
||||
# Decode connector will be able to create handshake with the prefill connector.
|
||||
decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
|
||||
decode_connector.register_kv_caches(kv_caches)
|
||||
|
||||
# Here we are testing the retrieval of NIXLAgentMetadata.
|
||||
# Knowing the implementation detail, we override the add_remote_agent
|
||||
@@ -402,6 +410,23 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
|
||||
self.kv_cache_layout = kv_cache_layout
|
||||
# Mock register_kv_caches attribute needed for tests that do not call it.
|
||||
self.src_xfer_handles_by_block_size = {self.block_size: 1}
|
||||
test_shape = self.attn_backend.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
self.kv_topo = TpKVTopology(
|
||||
tp_rank=self.tp_rank,
|
||||
engine_id=self.engine_id,
|
||||
remote_tp_size=self._tp_size, # shared state
|
||||
remote_block_size=self._block_size, # shared state
|
||||
is_mla=self.use_mla,
|
||||
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
|
||||
attn_backend=self.attn_backend,
|
||||
tensor_shape=test_shape,
|
||||
)
|
||||
|
||||
self.compat_hash = compute_nixl_compatibility_hash(
|
||||
self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
|
||||
)
|
||||
|
||||
def _nixl_handshake(
|
||||
self, host: str, port: int, remote_tp_size: int, expected_engine_id: str
|
||||
@@ -1370,6 +1395,7 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
|
||||
),
|
||||
),
|
||||
"TRITON_ATTN",
|
||||
"FLASHINFER",
|
||||
],
|
||||
)
|
||||
def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
@@ -1386,6 +1412,11 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
|
||||
vllm_config = create_vllm_config(attention_backend=attn_backend)
|
||||
|
||||
# Enable cross layers blocks
|
||||
vllm_config.kv_transfer_config.kv_connector_extra_config[
|
||||
"enable_cross_layers_blocks"
|
||||
] = True
|
||||
|
||||
# Import the appropriate backend based on the parameter
|
||||
if attn_backend == "FLASH_ATTN":
|
||||
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
@@ -1395,49 +1426,11 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
|
||||
|
||||
backend_cls = RocmAttentionBackend
|
||||
else: # TRITON_ATTN
|
||||
else: # TRITON
|
||||
from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
|
||||
|
||||
backend_cls = TritonAttentionBackend
|
||||
|
||||
# Create test kv cache tensors using proper backend shape
|
||||
kv_cache_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
|
||||
)
|
||||
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
kv_caches = {
|
||||
"layer0": shared_tensor,
|
||||
"layer1": unique_tensor,
|
||||
"layer2": shared_tensor,
|
||||
}
|
||||
|
||||
# Store tensor info for validation
|
||||
|
||||
test_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
|
||||
|
||||
if is_blocks_first:
|
||||
expected_tensor_size = shared_tensor.element_size() * shared_tensor.numel()
|
||||
expected_base_addrs = [
|
||||
shared_tensor.data_ptr(),
|
||||
unique_tensor.data_ptr(),
|
||||
]
|
||||
expected_num_entries = 2
|
||||
else:
|
||||
expected_tensor_size = (
|
||||
shared_tensor[0].element_size() * shared_tensor[0].numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
shared_tensor[0].data_ptr(),
|
||||
shared_tensor[1].data_ptr(),
|
||||
unique_tensor[0].data_ptr(),
|
||||
unique_tensor[1].data_ptr(),
|
||||
]
|
||||
expected_num_entries = 4
|
||||
|
||||
nixl_module = "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector"
|
||||
with (
|
||||
patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper,
|
||||
@@ -1466,6 +1459,107 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
# Reassure the shutdown() check that the thread is terminated
|
||||
mock_thread.return_value.is_alive.return_value = False
|
||||
|
||||
expected_tensor_size: int
|
||||
expected_base_addrs: list[int]
|
||||
expected_num_entries: int
|
||||
kv_caches: dict[str, torch.Tensor]
|
||||
if connector.prefer_cross_layer_blocks:
|
||||
num_layers = 32
|
||||
block_size = 16
|
||||
num_blocks = 8
|
||||
kv_cache_spec = AttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=4,
|
||||
head_size=64,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[
|
||||
KVCacheTensor(
|
||||
size=kv_cache_spec.page_size_bytes * num_blocks,
|
||||
shared_by=["dummy-layer"],
|
||||
)
|
||||
for i in range(num_layers)
|
||||
],
|
||||
# allocate_uniform_kv_caches does not use this
|
||||
kv_cache_groups=[],
|
||||
)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
_, cross_layers_kv_cache, _ = (
|
||||
KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
|
||||
kv_cache_config=kv_cache_config,
|
||||
attn_groups=[
|
||||
[
|
||||
AttentionGroup(
|
||||
backend=backend_cls,
|
||||
layer_names=[],
|
||||
kv_cache_spec=kv_cache_spec,
|
||||
kv_cache_group_id=0,
|
||||
)
|
||||
]
|
||||
],
|
||||
cache_dtype=torch.bfloat16,
|
||||
device=torch.cuda.current_device(),
|
||||
kernel_block_sizes=[block_size],
|
||||
)
|
||||
)
|
||||
# Store tensor info for validation
|
||||
expected_tensor_size = (
|
||||
cross_layers_kv_cache.element_size() * cross_layers_kv_cache.numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
cross_layers_kv_cache.data_ptr(),
|
||||
]
|
||||
expected_num_entries = 1
|
||||
|
||||
expected_blocks_count = 8
|
||||
|
||||
kv_caches = {"all-layers": cross_layers_kv_cache}
|
||||
|
||||
else:
|
||||
# Create test kv cache tensors using proper backend shape
|
||||
kv_cache_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
|
||||
)
|
||||
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
kv_caches = {
|
||||
"layer0": shared_tensor,
|
||||
"layer1": unique_tensor,
|
||||
"layer2": shared_tensor,
|
||||
}
|
||||
|
||||
# Store tensor info for validation
|
||||
|
||||
test_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
is_blocks_first = len(test_shape) == 5 and test_shape[0] == 1
|
||||
|
||||
if is_blocks_first:
|
||||
expected_tensor_size = (
|
||||
shared_tensor.element_size() * shared_tensor.numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
shared_tensor.data_ptr(),
|
||||
unique_tensor.data_ptr(),
|
||||
]
|
||||
expected_num_entries = 2
|
||||
else:
|
||||
expected_tensor_size = (
|
||||
shared_tensor[0].element_size() * shared_tensor[0].numel()
|
||||
)
|
||||
expected_base_addrs = [
|
||||
shared_tensor[0].data_ptr(),
|
||||
shared_tensor[1].data_ptr(),
|
||||
unique_tensor[0].data_ptr(),
|
||||
unique_tensor[1].data_ptr(),
|
||||
]
|
||||
expected_num_entries = 4
|
||||
expected_blocks_count = 8
|
||||
|
||||
# Execute register_kv_caches
|
||||
connector.register_kv_caches(kv_caches)
|
||||
|
||||
@@ -1489,16 +1583,19 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
|
||||
blocks_data, _ = mock_wrapper_instance.get_xfer_descs.call_args[0]
|
||||
|
||||
# Validate blocks_data structure and size
|
||||
expected_blocks_count = 8
|
||||
assert len(blocks_data) == expected_blocks_count, (
|
||||
f"Expected {expected_blocks_count} blocks, got {len(blocks_data)}"
|
||||
)
|
||||
|
||||
num_blocks = 2
|
||||
if is_blocks_first:
|
||||
expected_block_len = expected_tensor_size // num_blocks // 2
|
||||
else:
|
||||
if connector.prefer_cross_layer_blocks:
|
||||
num_blocks = 8
|
||||
expected_block_len = expected_tensor_size // num_blocks
|
||||
else:
|
||||
num_blocks = 2
|
||||
if is_blocks_first:
|
||||
expected_block_len = expected_tensor_size // num_blocks // 2
|
||||
else:
|
||||
expected_block_len = expected_tensor_size // num_blocks
|
||||
|
||||
for i, block_entry in enumerate(blocks_data):
|
||||
block_start_addr, block_len, tp_rank = block_entry
|
||||
@@ -2049,6 +2146,17 @@ def test_compatibility_hash_validation(
|
||||
)
|
||||
decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
|
||||
decode_worker = decode_connector.connector_worker
|
||||
kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
|
||||
num_blocks=2, block_size=16, num_kv_heads=4, head_size=64
|
||||
)
|
||||
shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)
|
||||
kv_caches = {
|
||||
"layer0": shared_tensor,
|
||||
"layer1": unique_tensor,
|
||||
"layer2": shared_tensor,
|
||||
}
|
||||
decode_connector.register_kv_caches(kv_caches)
|
||||
|
||||
remote_config_params: dict[str, Any] = {
|
||||
"model": "facebook/opt-125m",
|
||||
@@ -2071,7 +2179,9 @@ def test_compatibility_hash_validation(
|
||||
)
|
||||
)
|
||||
remote_hash = compute_nixl_compatibility_hash(
|
||||
remote_vllm_config, decode_worker.backend_name
|
||||
remote_vllm_config,
|
||||
decode_worker.backend_name,
|
||||
decode_worker.kv_topo.cross_layers_blocks,
|
||||
)
|
||||
|
||||
prefill_block_size = config_overrides.get("block_size", 16)
|
||||
@@ -2150,6 +2260,27 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
|
||||
decode_connector = NixlConnector(local_vllm_config, KVConnectorRole.WORKER)
|
||||
decode_worker = decode_connector.connector_worker
|
||||
|
||||
backend = get_current_attn_backend(local_vllm_config)
|
||||
test_shape = backend.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
decode_worker.kv_topo = TpKVTopology(
|
||||
tp_rank=decode_worker.tp_rank,
|
||||
engine_id=decode_worker.engine_id,
|
||||
remote_tp_size=decode_worker._tp_size, # shared state
|
||||
remote_block_size=decode_worker._block_size, # shared state
|
||||
is_mla=decode_worker.use_mla,
|
||||
total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
|
||||
attn_backend=backend,
|
||||
tensor_shape=test_shape,
|
||||
)
|
||||
|
||||
decode_worker.compat_hash = compute_nixl_compatibility_hash(
|
||||
decode_worker.vllm_config,
|
||||
decode_worker.backend_name,
|
||||
decode_worker.kv_topo.cross_layers_blocks,
|
||||
)
|
||||
|
||||
if error_scenario == "handshake_decode_error":
|
||||
msg_bytes = b"this is not valid msgpack data"
|
||||
elif error_scenario == "handshake_validation_error":
|
||||
|
||||
@@ -19,6 +19,7 @@ compressed-tensors, nm-testing/tinyllama-oneshot-w8a16-per-channel, main
|
||||
compressed-tensors, nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test, main
|
||||
compressed-tensors, nm-testing/Phi-3-mini-128k-instruct-FP8, main
|
||||
compressed-tensors, neuralmagic/Phi-3-medium-128k-instruct-quantized.w4a16, main
|
||||
compressed-tensors, nm-testing/TinyLlama-1.1B-Chat-v1.0-actorder-group, main
|
||||
#compressed-tensors, mgoin/DeepSeek-Coder-V2-Lite-Instruct-FP8, main
|
||||
compressed-tensors, nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-FP8-Dynamic-testing, main, 90
|
||||
compressed-tensors, nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-W8A8-testing, main, 90
|
||||
|
||||
@@ -1,233 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Generate S3 PyPI Root Index for Latest Version
|
||||
#
|
||||
# Creates a PEP 503 compatible index.html at rocm/ pointing to the latest
|
||||
# semantic version's packages. This enables users to install with:
|
||||
# uv pip install vllm --extra-index-url s3://vllm-wheels/rocm
|
||||
#
|
||||
# Usage:
|
||||
# generate-root-index.sh [options]
|
||||
#
|
||||
# Options:
|
||||
# --dry-run Preview changes without uploading
|
||||
# --version VER Use specific version instead of auto-detecting latest
|
||||
#
|
||||
# Environment variables:
|
||||
# S3_BUCKET - Bucket name (default: vllm-wheels)
|
||||
# VARIANT - ROCm variant (default: rocm700)
|
||||
# DRY_RUN - Set to 1 for preview mode (same as --dry-run)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# ======== Configuration ========
|
||||
BUCKET="${S3_BUCKET:-vllm-wheels}"
|
||||
VARIANT="${VARIANT:-rocm700}"
|
||||
DRY_RUN="${DRY_RUN:-0}"
|
||||
FORCE_VERSION=""
|
||||
|
||||
# Parse command line arguments
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
--dry-run)
|
||||
DRY_RUN=1
|
||||
shift
|
||||
;;
|
||||
--version)
|
||||
FORCE_VERSION="$2"
|
||||
shift 2
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# Working directory for generated files
|
||||
WORK_DIR=$(mktemp -d)
|
||||
trap 'rm -rf "$WORK_DIR"' EXIT
|
||||
|
||||
echo "========================================"
|
||||
echo "Generate Root Index for Latest Version"
|
||||
echo "========================================"
|
||||
echo "S3 Bucket: $BUCKET"
|
||||
echo "ROCm Variant: $VARIANT"
|
||||
echo "Dry Run: $DRY_RUN"
|
||||
echo "========================================"
|
||||
echo ""
|
||||
|
||||
# ======== Step 1: Find latest semantic version ========
|
||||
|
||||
echo "Step 1: Finding latest semantic version..."
|
||||
|
||||
# List all directories under rocm/
|
||||
aws s3api list-objects-v2 \
|
||||
--bucket "$BUCKET" \
|
||||
--prefix "rocm/" \
|
||||
--delimiter "/" \
|
||||
--query 'CommonPrefixes[].Prefix' \
|
||||
--output text | tr '\t' '\n' > "$WORK_DIR/all_prefixes.txt"
|
||||
|
||||
# Filter for semantic versions (x.y.z pattern)
|
||||
grep -oE 'rocm/[0-9]+\.[0-9]+\.[0-9]+/' "$WORK_DIR/all_prefixes.txt" | \
|
||||
sed 's|rocm/||; s|/||' | \
|
||||
sort -V > "$WORK_DIR/versions.txt" || true
|
||||
|
||||
if [[ ! -s "$WORK_DIR/versions.txt" ]]; then
|
||||
echo "ERROR: No semantic versions found under s3://$BUCKET/rocm/"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Found versions:"
|
||||
cat "$WORK_DIR/versions.txt"
|
||||
echo ""
|
||||
|
||||
if [[ -n "$FORCE_VERSION" ]]; then
|
||||
LATEST_VERSION="$FORCE_VERSION"
|
||||
echo "Using forced version: $LATEST_VERSION"
|
||||
else
|
||||
LATEST_VERSION=$(tail -1 "$WORK_DIR/versions.txt")
|
||||
echo "Latest version (auto-detected): $LATEST_VERSION"
|
||||
fi
|
||||
|
||||
# Verify the version exists
|
||||
if ! grep -qx "$LATEST_VERSION" "$WORK_DIR/versions.txt"; then
|
||||
echo "ERROR: Version $LATEST_VERSION not found in bucket"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# ======== Step 2: List packages from latest version ========
|
||||
|
||||
echo ""
|
||||
echo "Step 2: Listing packages from rocm/$LATEST_VERSION/$VARIANT/..."
|
||||
|
||||
VERSION_PREFIX="rocm/$LATEST_VERSION/$VARIANT/"
|
||||
|
||||
# List package directories
|
||||
aws s3api list-objects-v2 \
|
||||
--bucket "$BUCKET" \
|
||||
--prefix "$VERSION_PREFIX" \
|
||||
--delimiter "/" \
|
||||
--query 'CommonPrefixes[].Prefix' \
|
||||
--output text | tr '\t' '\n' > "$WORK_DIR/package_prefixes.txt" || true
|
||||
|
||||
if [[ ! -s "$WORK_DIR/package_prefixes.txt" ]]; then
|
||||
echo "ERROR: No packages found under s3://$BUCKET/$VERSION_PREFIX"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Extract package names
|
||||
sed "s|${VERSION_PREFIX}||; s|/||g" "$WORK_DIR/package_prefixes.txt" | \
|
||||
grep -v '^$' > "$WORK_DIR/packages.txt"
|
||||
|
||||
echo "Found packages:"
|
||||
cat "$WORK_DIR/packages.txt"
|
||||
echo ""
|
||||
|
||||
# ======== Step 3: Generate root index.html ========
|
||||
|
||||
echo "Step 3: Generating root index.html..."
|
||||
|
||||
mkdir -p "$WORK_DIR/output"
|
||||
|
||||
{
|
||||
cat <<'EOF'
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<meta name="pypi:repository-version" content="1.0">
|
||||
</head>
|
||||
<body>
|
||||
EOF
|
||||
|
||||
while read -r pkg; do
|
||||
echo " <a href=\"$pkg/\">$pkg</a><br>"
|
||||
done < "$WORK_DIR/packages.txt"
|
||||
|
||||
cat <<'EOF'
|
||||
</body>
|
||||
</html>
|
||||
EOF
|
||||
} > "$WORK_DIR/output/index.html"
|
||||
|
||||
echo "Generated root index.html:"
|
||||
cat "$WORK_DIR/output/index.html"
|
||||
echo ""
|
||||
|
||||
# ======== Step 4: Copy and adjust package index files ========
|
||||
|
||||
echo "Step 4: Copying and adjusting package index files..."
|
||||
|
||||
while read -r pkg; do
|
||||
echo "Processing package: $pkg"
|
||||
|
||||
# Download existing index.html from versioned path
|
||||
SOURCE_INDEX="s3://$BUCKET/$VERSION_PREFIX$pkg/index.html"
|
||||
|
||||
mkdir -p "$WORK_DIR/output/$pkg"
|
||||
|
||||
if aws s3 cp "$SOURCE_INDEX" "$WORK_DIR/output/$pkg/index.html" 2>/dev/null; then
|
||||
# Adjust relative paths:
|
||||
# Original: href="../../../{commit}/wheel.whl" (from rocm/0.13.0/rocm710/vllm/)
|
||||
# New: href="../{commit}/wheel.whl" (from rocm/vllm/)
|
||||
sed -i 's|href="\.\./\.\./\.\./|href="../|g' "$WORK_DIR/output/$pkg/index.html"
|
||||
echo " - Downloaded and adjusted: $pkg/index.html"
|
||||
else
|
||||
echo " - WARNING: Could not download index for $pkg"
|
||||
fi
|
||||
done < "$WORK_DIR/packages.txt"
|
||||
|
||||
echo ""
|
||||
|
||||
# ======== Step 5: Upload to S3 ========
|
||||
|
||||
echo "Step 5: Uploading to s3://$BUCKET/rocm/..."
|
||||
echo ""
|
||||
|
||||
# List what would be uploaded
|
||||
echo "Files to upload:"
|
||||
find "$WORK_DIR/output" -name "*.html" -type f | while read -r file; do
|
||||
rel_path="${file#$WORK_DIR/output/}"
|
||||
echo " rocm/$rel_path"
|
||||
done
|
||||
echo ""
|
||||
|
||||
if [[ "$DRY_RUN" == "1" ]]; then
|
||||
echo "DRY RUN - Skipping upload"
|
||||
echo ""
|
||||
echo "Preview of generated files:"
|
||||
echo "----------------------------------------"
|
||||
echo "rocm/index.html:"
|
||||
cat "$WORK_DIR/output/index.html"
|
||||
echo ""
|
||||
echo "----------------------------------------"
|
||||
echo "Sample package index (first package):"
|
||||
FIRST_PKG=$(head -1 "$WORK_DIR/packages.txt")
|
||||
if [[ -f "$WORK_DIR/output/$FIRST_PKG/index.html" ]]; then
|
||||
echo "rocm/$FIRST_PKG/index.html:"
|
||||
cat "$WORK_DIR/output/$FIRST_PKG/index.html"
|
||||
fi
|
||||
else
|
||||
# Upload all generated files
|
||||
aws s3 cp --recursive "$WORK_DIR/output/" "s3://$BUCKET/rocm/" \
|
||||
--content-type "text/html"
|
||||
|
||||
echo "Upload complete!"
|
||||
fi
|
||||
|
||||
# ======== Summary ========
|
||||
|
||||
echo ""
|
||||
echo "========================================"
|
||||
echo "Root Index Generation Complete!"
|
||||
echo "========================================"
|
||||
echo ""
|
||||
echo "Latest version: $LATEST_VERSION"
|
||||
echo "Packages indexed: $(wc -l < "$WORK_DIR/packages.txt")"
|
||||
echo ""
|
||||
echo "Install command:"
|
||||
echo " uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/"
|
||||
echo "========================================"
|
||||
+42
-11
@@ -900,8 +900,6 @@ def cutlass_sparse_scaled_mm_supported(cuda_device_capability: int) -> bool:
|
||||
|
||||
|
||||
def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool:
|
||||
if cuda_device_capability < 90 or cuda_device_capability >= 110:
|
||||
return False
|
||||
try:
|
||||
return torch.ops._C.cutlass_group_gemm_supported(cuda_device_capability)
|
||||
except AttributeError:
|
||||
@@ -2034,20 +2032,35 @@ def selective_scan_fwd(
|
||||
)
|
||||
|
||||
|
||||
# NOTE: The wvSplitK kernel (and all of the kernels in skinny_gemms.cu)
|
||||
# are unable to properly handle non-contiguous
|
||||
# tensors. It might be a good TODO(rasmith) to augment these kernels
|
||||
# to be able to handle non-contiguous kernels for better performance.
|
||||
def rocm_enforce_contiguous_skinny_gemm_inputs(
|
||||
a: torch.Tensor, b: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
a = a.contiguous() # no-op if already contiguous, else clone
|
||||
b = b.contiguous() # no-op if already contiguous, else clone
|
||||
return a, b
|
||||
|
||||
|
||||
# ROCm skinny gemms
|
||||
def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor:
|
||||
a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
|
||||
return torch.ops._rocm_C.LLMM1(a, b, rows_per_block)
|
||||
|
||||
|
||||
def wvSplitK(
|
||||
a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
|
||||
) -> torch.Tensor:
|
||||
a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
|
||||
return torch.ops._rocm_C.wvSplitK(a, b, bias, cu_count)
|
||||
|
||||
|
||||
def wvSplitKrc(
|
||||
a: torch.Tensor, b: torch.Tensor, cu_count: int, bias: torch.Tensor = None
|
||||
) -> torch.Tensor:
|
||||
a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
|
||||
return torch.ops._rocm_C.wvSplitKrc(a, b, bias, cu_count)
|
||||
|
||||
|
||||
@@ -2060,6 +2073,7 @@ def wvSplitKQ(
|
||||
cu_count: int,
|
||||
bias: torch.Tensor = None,
|
||||
) -> torch.Tensor:
|
||||
a, b = rocm_enforce_contiguous_skinny_gemm_inputs(a, b)
|
||||
out = torch.empty((b.shape[0], a.shape[0]), dtype=out_dtype, device=b.device)
|
||||
torch.ops._rocm_C.wvSplitKQ(a, b, bias, out, scale_a, scale_b, cu_count)
|
||||
return out
|
||||
@@ -2831,13 +2845,13 @@ if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"):
|
||||
|
||||
class CPUDNNLGEMMHandler:
|
||||
def __init__(self) -> None:
|
||||
self.handler: int | None = None
|
||||
self.handler_tensor: torch.Tensor | None = None
|
||||
self.n = -1
|
||||
self.k = -1
|
||||
|
||||
def __del__(self):
|
||||
if self.handler is not None:
|
||||
torch.ops._C.release_dnnl_matmul_handler(self.handler)
|
||||
if self.handler_tensor is not None:
|
||||
torch.ops._C.release_dnnl_matmul_handler(self.handler_tensor.item())
|
||||
|
||||
|
||||
_supports_onednn = bool(hasattr(torch.ops._C, "create_onednn_mm_handler"))
|
||||
@@ -2853,8 +2867,10 @@ def create_onednn_mm(
|
||||
) -> CPUDNNLGEMMHandler:
|
||||
handler = CPUDNNLGEMMHandler()
|
||||
handler.k, handler.n = weight.size()
|
||||
handler.handler = torch.ops._C.create_onednn_mm_handler(
|
||||
weight, primitive_cache_size
|
||||
# store the handler pointer in a tensor it doesn't get inlined
|
||||
handler.handler_tensor = torch.tensor(
|
||||
torch.ops._C.create_onednn_mm_handler(weight, primitive_cache_size),
|
||||
dtype=torch.int64,
|
||||
)
|
||||
return handler
|
||||
|
||||
@@ -2866,7 +2882,7 @@ def onednn_mm(
|
||||
) -> torch.Tensor:
|
||||
output = torch.empty((*x.shape[0:-1], dnnl_handler.n), dtype=x.dtype)
|
||||
torch.ops._C.onednn_mm(
|
||||
output, x.reshape(-1, dnnl_handler.k), bias, dnnl_handler.handler
|
||||
output, x.reshape(-1, dnnl_handler.k), bias, dnnl_handler.handler_tensor
|
||||
)
|
||||
|
||||
return output
|
||||
@@ -2882,8 +2898,17 @@ def create_onednn_scaled_mm(
|
||||
) -> CPUDNNLGEMMHandler:
|
||||
handler = CPUDNNLGEMMHandler()
|
||||
handler.k, handler.n = weight.size()
|
||||
handler.handler = torch.ops._C.create_onednn_scaled_mm_handler(
|
||||
weight, weight_scales, output_type, dynamic_quant, use_azp, primitive_cache_size
|
||||
# store the handler pointer in a tensor so it doesn't get inlined
|
||||
handler.handler_tensor = torch.tensor(
|
||||
torch.ops._C.create_onednn_scaled_mm_handler(
|
||||
weight,
|
||||
weight_scales,
|
||||
output_type,
|
||||
dynamic_quant,
|
||||
use_azp,
|
||||
primitive_cache_size,
|
||||
),
|
||||
dtype=torch.int64,
|
||||
)
|
||||
return handler
|
||||
|
||||
@@ -2936,7 +2961,13 @@ def onednn_scaled_mm(
|
||||
bias: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
torch.ops._C.onednn_scaled_mm(
|
||||
output, x, input_scale, input_zp, input_zp_adj, bias, dnnl_handler.handler
|
||||
output,
|
||||
x,
|
||||
input_scale,
|
||||
input_zp,
|
||||
input_zp_adj,
|
||||
bias,
|
||||
dnnl_handler.handler_tensor,
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
@@ -361,14 +361,7 @@ def split_graph(
|
||||
subgraph_id += 1
|
||||
node_to_subgraph_id[node] = subgraph_id
|
||||
split_op_graphs.append(subgraph_id)
|
||||
|
||||
# keep consecutive splitting ops together
|
||||
# (we know node.next exists because node isn't the last (output) node)
|
||||
if should_split(node.next, splitting_ops):
|
||||
# this will get incremented by the next node
|
||||
subgraph_id -= 1
|
||||
else:
|
||||
subgraph_id += 1
|
||||
subgraph_id += 1
|
||||
else:
|
||||
node_to_subgraph_id[node] = subgraph_id
|
||||
|
||||
|
||||
@@ -280,9 +280,10 @@ class DynamicShapesConfig:
|
||||
until this change picked up https://github.com/pytorch/pytorch/pull/169239.
|
||||
"""
|
||||
|
||||
assume_32_bit_indexing: bool = True
|
||||
assume_32_bit_indexing: bool = False
|
||||
"""
|
||||
whether all tensor sizes can use 32 bit indexing.
|
||||
`True` requires PyTorch 2.10+
|
||||
"""
|
||||
|
||||
def compute_hash(self) -> str:
|
||||
@@ -581,24 +582,6 @@ class CompilationConfig:
|
||||
local_cache_dir: str = field(default=None, init=False) # type: ignore
|
||||
"""local cache dir for each rank"""
|
||||
|
||||
fast_moe_cold_start = True
|
||||
"""Optimization for fast MOE cold start.
|
||||
|
||||
This is a bit of a hack that assumes that:
|
||||
1. the only decoder forward pass being run is the current model
|
||||
2. the decoder forward pass runs all of the MOEs in the order in which they
|
||||
are initialized
|
||||
|
||||
When the above two conditions hold, this option greatly decreases cold start
|
||||
time for MOE models.
|
||||
|
||||
If the above two conditions don't hold, then this option will lead to silent
|
||||
incorrectness. The only condition in which this doesn't hold is speculative
|
||||
decoding, where there is a draft model that may have MOEs in them.
|
||||
|
||||
NB: We're working on a longer-term solution that doesn't need these assumptions.
|
||||
"""
|
||||
|
||||
# keep track of enabled and disabled custom ops
|
||||
enabled_custom_ops: Counter[str] = field(default_factory=Counter, init=False)
|
||||
"""custom ops that are enabled"""
|
||||
@@ -614,10 +597,6 @@ class CompilationConfig:
|
||||
Map from layer name to layer objects that need to be accessed outside
|
||||
model code, e.g., Attention, FusedMOE when dp_size>1."""
|
||||
|
||||
static_all_moe_layers: list[str] = field(default_factory=list, init=False)
|
||||
"""The names of all the MOE layers in the model
|
||||
"""
|
||||
|
||||
# Attention ops; used for piecewise cudagraphs
|
||||
# Use PyTorch operator format: "namespace::name"
|
||||
_attention_ops: ClassVar[list[str]] = [
|
||||
@@ -947,15 +926,6 @@ class CompilationConfig:
|
||||
# for details. Make a copy to avoid mutating the class-level
|
||||
# list via reference.
|
||||
self.splitting_ops = list(self._attention_ops)
|
||||
|
||||
# unified_kv_cache_update has a string param that prevents Inductor
|
||||
# from reusing piecewise graphs. Remove it from the compiled graph.
|
||||
# This has the side-effect of excluding cache from cudagraphs but
|
||||
# that doesn't seem to affect performance.
|
||||
# https://github.com/vllm-project/vllm/issues/33267
|
||||
if not self.use_inductor_graph_partition:
|
||||
self.splitting_ops.append("vllm::unified_kv_cache_update")
|
||||
|
||||
elif len(self.splitting_ops) == 0:
|
||||
if (
|
||||
self.cudagraph_mode == CUDAGraphMode.PIECEWISE
|
||||
|
||||
@@ -34,13 +34,13 @@ MTPModelTypes = Literal[
|
||||
"mimo_mtp",
|
||||
"glm4_moe_mtp",
|
||||
"glm4_moe_lite_mtp",
|
||||
"glm_ocr_mtp",
|
||||
"ernie_mtp",
|
||||
"exaone_moe_mtp",
|
||||
"qwen3_next_mtp",
|
||||
"longcat_flash_mtp",
|
||||
"mtp",
|
||||
"pangu_ultra_moe_mtp",
|
||||
"step3p5_mtp",
|
||||
]
|
||||
EagleModelTypes = Literal["eagle", "eagle3", MTPModelTypes]
|
||||
SpeculativeMethod = Literal[
|
||||
@@ -222,6 +222,17 @@ class SpeculativeConfig:
|
||||
}
|
||||
)
|
||||
|
||||
if hf_config.architectures[0] == "GlmOcrForConditionalGeneration":
|
||||
hf_config.model_type = "glm_ocr_mtp"
|
||||
n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
|
||||
hf_config.update(
|
||||
{
|
||||
"num_hidden_layers": 0,
|
||||
"n_predict": n_predict,
|
||||
"architectures": ["GlmOcrMTPModel"],
|
||||
}
|
||||
)
|
||||
|
||||
if hf_config.model_type == "ernie4_5_moe":
|
||||
hf_config.model_type = "ernie_mtp"
|
||||
if hf_config.model_type == "ernie_mtp":
|
||||
@@ -253,11 +264,6 @@ class SpeculativeConfig:
|
||||
{"n_predict": n_predict, "architectures": ["LongCatFlashMTPModel"]}
|
||||
)
|
||||
|
||||
if hf_config.model_type == "step3p5":
|
||||
hf_config.model_type = "step3p5_mtp"
|
||||
n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)
|
||||
hf_config.update({"n_predict": n_predict, "architectures": ["Step3p5MTP"]})
|
||||
|
||||
if initial_architecture == "MistralLarge3ForCausalLM":
|
||||
hf_config.update({"architectures": ["EagleMistralLarge3ForCausalLM"]})
|
||||
|
||||
|
||||
@@ -59,7 +59,7 @@ class NaiveAll2AllManager(All2AllManagerBase):
|
||||
|
||||
return buffer
|
||||
|
||||
def dispatch(
|
||||
def dispatch_router_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
@@ -84,6 +84,34 @@ class NaiveAll2AllManager(All2AllManagerBase):
|
||||
|
||||
return hidden_states, router_logits
|
||||
|
||||
def dispatch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
if extra_tensors is not None:
|
||||
raise NotImplementedError(
|
||||
"extra_tensors is not supported for NaiveAll2AllManager"
|
||||
)
|
||||
sp_size = self.tp_group.world_size if is_sequence_parallel else 1
|
||||
dp_metadata = get_forward_context().dp_metadata
|
||||
assert dp_metadata is not None
|
||||
cu_tokens_across_sp_cpu = dp_metadata.cu_tokens_across_sp(sp_size)
|
||||
|
||||
hidden_states = self.naive_multicast(
|
||||
hidden_states, cu_tokens_across_sp_cpu, is_sequence_parallel
|
||||
)
|
||||
topk_weights = self.naive_multicast(
|
||||
topk_weights, cu_tokens_across_sp_cpu, is_sequence_parallel
|
||||
)
|
||||
topk_ids = self.naive_multicast(
|
||||
topk_ids, cu_tokens_across_sp_cpu, is_sequence_parallel
|
||||
)
|
||||
return hidden_states, topk_weights, topk_ids
|
||||
|
||||
def combine(
|
||||
self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
|
||||
) -> torch.Tensor:
|
||||
@@ -114,7 +142,7 @@ class AgRsAll2AllManager(All2AllManagerBase):
|
||||
def __init__(self, cpu_group):
|
||||
super().__init__(cpu_group)
|
||||
|
||||
def dispatch(
|
||||
def dispatch_router_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
@@ -148,6 +176,46 @@ class AgRsAll2AllManager(All2AllManagerBase):
|
||||
return (gathered_tensors[0], gathered_tensors[1], gathered_tensors[2:])
|
||||
return gathered_tensors[0], gathered_tensors[1]
|
||||
|
||||
def dispatch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
"""
|
||||
Gather hidden_states and router_logits from all dp ranks.
|
||||
"""
|
||||
dp_metadata = get_forward_context().dp_metadata
|
||||
assert dp_metadata is not None
|
||||
sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
|
||||
assert sizes is not None
|
||||
dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
|
||||
assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
|
||||
|
||||
tensors_to_gather = [hidden_states, topk_weights, topk_ids]
|
||||
if extra_tensors is not None:
|
||||
tensors_to_gather.extend(extra_tensors)
|
||||
|
||||
gathered_tensors = dist_group.all_gatherv(
|
||||
tensors_to_gather,
|
||||
dim=0,
|
||||
sizes=sizes,
|
||||
)
|
||||
|
||||
hidden_states = gathered_tensors[0]
|
||||
topk_weights = gathered_tensors[1]
|
||||
topk_ids = gathered_tensors[2]
|
||||
|
||||
if extra_tensors is None:
|
||||
return hidden_states, topk_weights, topk_ids
|
||||
|
||||
return hidden_states, topk_weights, topk_ids, gathered_tensors[3:]
|
||||
|
||||
def combine(
|
||||
self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
|
||||
) -> torch.Tensor:
|
||||
@@ -216,7 +284,7 @@ class PPLXAll2AllManager(All2AllManagerBase):
|
||||
pplx.AllToAll.internode if self.internode else pplx.AllToAll.intranode,
|
||||
)
|
||||
|
||||
def dispatch(
|
||||
def dispatch_router_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
@@ -225,6 +293,19 @@ class PPLXAll2AllManager(All2AllManagerBase):
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
raise NotImplementedError
|
||||
|
||||
def dispatch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
raise NotImplementedError
|
||||
|
||||
def combine(
|
||||
self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
|
||||
) -> torch.Tensor:
|
||||
@@ -264,7 +345,7 @@ class DeepEPAll2AllManagerBase(All2AllManagerBase):
|
||||
def get_handle(self, kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
def dispatch(
|
||||
def dispatch_router_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
@@ -273,6 +354,19 @@ class DeepEPAll2AllManagerBase(All2AllManagerBase):
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
raise NotImplementedError
|
||||
|
||||
def dispatch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
raise NotImplementedError
|
||||
|
||||
def combine(
|
||||
self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
|
||||
) -> torch.Tensor:
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import threading
|
||||
from typing import Any
|
||||
from weakref import WeakValueDictionary
|
||||
|
||||
import torch
|
||||
@@ -64,13 +63,32 @@ class All2AllManagerBase:
|
||||
# and reuse it for the same config.
|
||||
raise NotImplementedError
|
||||
|
||||
def dispatch(
|
||||
def dispatch_router_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> Any:
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
# Subclasses should either:
|
||||
# - implement handling for extra_tensors, or
|
||||
# - raise a clear error if extra_tensors is not supported.
|
||||
raise NotImplementedError
|
||||
|
||||
def dispatch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
# Subclasses should either:
|
||||
# - implement handling for extra_tensors, or
|
||||
# - raise a clear error if extra_tensors is not supported.
|
||||
@@ -280,7 +298,7 @@ class DeviceCommunicatorBase:
|
||||
for module in moe_modules:
|
||||
module.maybe_init_modular_kernel()
|
||||
|
||||
def dispatch(
|
||||
def dispatch_router_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
@@ -294,8 +312,29 @@ class DeviceCommunicatorBase:
|
||||
Dispatch the hidden states and router logits to the appropriate device.
|
||||
This is a no-op in the base class.
|
||||
"""
|
||||
if extra_tensors is not None:
|
||||
return hidden_states, router_logits, extra_tensors
|
||||
return hidden_states, router_logits
|
||||
|
||||
def dispatch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
"""
|
||||
Dispatch the hidden states and topk weights/ids to the appropriate device.
|
||||
This is a no-op in the base class.
|
||||
"""
|
||||
if extra_tensors is not None:
|
||||
return hidden_states, topk_weights, topk_ids, extra_tensors
|
||||
return hidden_states, topk_weights, topk_ids
|
||||
|
||||
def combine(
|
||||
self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
|
||||
) -> torch.Tensor:
|
||||
|
||||
@@ -130,29 +130,65 @@ class CpuCommunicator(DeviceCommunicatorBase):
|
||||
) -> dict[str, torch.Tensor | Any]:
|
||||
return self.dist_module.recv_tensor_dict(src)
|
||||
|
||||
def dispatch( # type: ignore[override]
|
||||
def dispatch_router_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
"""
|
||||
Dispatch the hidden states and router logits to the appropriate device.
|
||||
This is a no-op in the base class.
|
||||
"""
|
||||
|
||||
assert self.all2all_manager is not None
|
||||
return self.all2all_manager.dispatch(
|
||||
return self.all2all_manager.dispatch_router_logits(
|
||||
hidden_states,
|
||||
router_logits,
|
||||
is_sequence_parallel,
|
||||
extra_tensors, # type: ignore[call-arg]
|
||||
extra_tensors,
|
||||
)
|
||||
|
||||
def dispatch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
"""
|
||||
Dispatch the hidden states and topk weights/ids to the appropriate device.
|
||||
This is a no-op in the base class.
|
||||
"""
|
||||
assert self.all2all_manager is not None
|
||||
return self.all2all_manager.dispatch(
|
||||
hidden_states,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
is_sequence_parallel,
|
||||
extra_tensors=extra_tensors,
|
||||
)
|
||||
|
||||
def combine(
|
||||
self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Combine the hidden states and router logits from the appropriate device.
|
||||
This is a no-op in the base class.
|
||||
"""
|
||||
assert self.all2all_manager is not None
|
||||
hidden_states = self.all2all_manager.combine(
|
||||
hidden_states, is_sequence_parallel
|
||||
return self.all2all_manager.combine(
|
||||
hidden_states,
|
||||
is_sequence_parallel,
|
||||
)
|
||||
return hidden_states
|
||||
|
||||
|
||||
class _CPUSHMDistributed:
|
||||
|
||||
@@ -322,7 +322,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
|
||||
|
||||
return output_list
|
||||
|
||||
def dispatch( # type: ignore[override]
|
||||
def dispatch_router_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
@@ -332,19 +332,52 @@ class CudaCommunicator(DeviceCommunicatorBase):
|
||||
tuple[torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
"""
|
||||
Dispatch the hidden states and router logits to the appropriate device.
|
||||
This is a no-op in the base class.
|
||||
"""
|
||||
|
||||
assert self.all2all_manager is not None
|
||||
return self.all2all_manager.dispatch(
|
||||
return self.all2all_manager.dispatch_router_logits(
|
||||
hidden_states,
|
||||
router_logits,
|
||||
is_sequence_parallel,
|
||||
extra_tensors, # type: ignore[call-arg]
|
||||
extra_tensors,
|
||||
)
|
||||
|
||||
def dispatch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
"""
|
||||
Dispatch the hidden states and topk weights/ids to the appropriate device.
|
||||
This is a no-op in the base class.
|
||||
"""
|
||||
assert self.all2all_manager is not None
|
||||
return self.all2all_manager.dispatch(
|
||||
hidden_states,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
is_sequence_parallel,
|
||||
extra_tensors=extra_tensors,
|
||||
)
|
||||
|
||||
def combine(
|
||||
self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Combine the hidden states and router logits from the appropriate device.
|
||||
This is a no-op in the base class.
|
||||
"""
|
||||
assert self.all2all_manager is not None
|
||||
hidden_states = self.all2all_manager.combine(
|
||||
hidden_states, is_sequence_parallel
|
||||
return self.all2all_manager.combine(
|
||||
hidden_states,
|
||||
is_sequence_parallel,
|
||||
)
|
||||
return hidden_states
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from typing import Any
|
||||
|
||||
import torch.distributed as dist
|
||||
from flashinfer.comm.mnnvl import CommBackend as CommBackend
|
||||
|
||||
@@ -23,5 +25,14 @@ class CustomCommunicator(CommBackend):
|
||||
dist.all_gather_object(gathered, data, group=self._group)
|
||||
return gathered
|
||||
|
||||
# NOTE(rob): CommBackend is an abstract class, and bcast/barrier
|
||||
# are unimplemented on vLLM side. If we need to utilize these
|
||||
# methods in the future, can create a concrete implementation.
|
||||
def bcast(self, data: Any, root: int) -> Any:
|
||||
raise NotImplementedError
|
||||
|
||||
def barrier(self) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def Split(self, color: int, key: int) -> "CustomCommunicator":
|
||||
return self
|
||||
|
||||
@@ -196,26 +196,62 @@ class XpuCommunicator(DeviceCommunicatorBase):
|
||||
def broadcast(self, input_: torch.Tensor, src: int = 0) -> None:
|
||||
dist.broadcast(input_, src=src, group=self.device_group)
|
||||
|
||||
def dispatch(
|
||||
def dispatch_router_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
"""
|
||||
Dispatch the hidden states and router logits to the appropriate device.
|
||||
This is a no-op in the base class.
|
||||
"""
|
||||
|
||||
assert self.all2all_manager is not None
|
||||
return self.all2all_manager.dispatch(
|
||||
return self.all2all_manager.dispatch_router_logits(
|
||||
hidden_states,
|
||||
router_logits,
|
||||
is_sequence_parallel,
|
||||
extra_tensors, # type: ignore[call-arg]
|
||||
extra_tensors,
|
||||
)
|
||||
|
||||
def dispatch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor, torch.Tensor]
|
||||
| tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
"""
|
||||
Dispatch the hidden states and topk weights/ids to the appropriate device.
|
||||
This is a no-op in the base class.
|
||||
"""
|
||||
assert self.all2all_manager is not None
|
||||
return self.all2all_manager.dispatch(
|
||||
hidden_states,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
is_sequence_parallel,
|
||||
extra_tensors=extra_tensors,
|
||||
)
|
||||
|
||||
def combine(
|
||||
self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Combine the hidden states and router logits from the appropriate device.
|
||||
This is a no-op in the base class.
|
||||
"""
|
||||
assert self.all2all_manager is not None
|
||||
hidden_states = self.all2all_manager.combine(
|
||||
hidden_states, is_sequence_parallel
|
||||
return self.all2all_manager.combine(
|
||||
hidden_states,
|
||||
is_sequence_parallel,
|
||||
)
|
||||
return hidden_states
|
||||
|
||||
@@ -316,12 +316,13 @@ class TpKVTopology:
|
||||
attn_backend: type[AttentionBackend]
|
||||
engine_id: EngineId
|
||||
remote_block_size: dict[EngineId, int]
|
||||
tensor_shape: torch.Size | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
# Figure out whether the first dimension of the cache is K/V
|
||||
# or num_blocks. This is used to register the memory regions correctly.
|
||||
kv_cache_shape = self.attn_backend.get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
num_blocks=1, block_size=16, num_kv_heads=4, head_size=1
|
||||
)
|
||||
# Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D],
|
||||
# we just mock num_blocks to 1 for the dimension check below.
|
||||
@@ -329,6 +330,32 @@ class TpKVTopology:
|
||||
len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1
|
||||
)
|
||||
|
||||
self._kv_heads_position: int | None = None
|
||||
self._cross_layers_blocks = False
|
||||
if self.tensor_shape is not None:
|
||||
self._cross_layers_blocks = (
|
||||
len(self.tensor_shape) == len(kv_cache_shape) + 1
|
||||
)
|
||||
|
||||
if self._cross_layers_blocks:
|
||||
# prepend layers dimension
|
||||
kv_cache_shape = (80,) + kv_cache_shape
|
||||
try:
|
||||
kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order(
|
||||
include_num_layers_dimension=self._cross_layers_blocks
|
||||
)
|
||||
except (AttributeError, NotImplementedError):
|
||||
kv_cache_stride_order = tuple(range(len(self.tensor_shape)))
|
||||
|
||||
# permute kv_cache_shape according to stride_order
|
||||
kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
|
||||
|
||||
physical_block_size_position = kv_cache_shape.index(16)
|
||||
assert physical_block_size_position is not None
|
||||
self._physical_block_size_position = -(
|
||||
len(kv_cache_shape) - physical_block_size_position
|
||||
)
|
||||
|
||||
@property
|
||||
def is_kv_layout_blocks_first(self) -> bool:
|
||||
return self._is_kv_layout_blocks_first
|
||||
@@ -336,7 +363,9 @@ class TpKVTopology:
|
||||
@property
|
||||
def split_k_and_v(self) -> bool:
|
||||
# Whether to register regions for K and V separately (when present).
|
||||
return not (self.is_mla or self.is_kv_layout_blocks_first)
|
||||
return not (
|
||||
self._cross_layers_blocks or self.is_mla or self.is_kv_layout_blocks_first
|
||||
)
|
||||
|
||||
@property
|
||||
def tp_size(self) -> int:
|
||||
@@ -346,6 +375,14 @@ class TpKVTopology:
|
||||
def block_size(self) -> int:
|
||||
return self.remote_block_size[self.engine_id]
|
||||
|
||||
@property
|
||||
def cross_layers_blocks(self) -> bool:
|
||||
return self._cross_layers_blocks
|
||||
|
||||
@property
|
||||
def block_size_position(self) -> int:
|
||||
return self._physical_block_size_position
|
||||
|
||||
def tp_ratio(
|
||||
self,
|
||||
remote_tp_size: int,
|
||||
|
||||
@@ -54,7 +54,7 @@ from vllm.forward_context import ForwardContext
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
|
||||
from vllm.v1.attention.backend import AttentionMetadata
|
||||
from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
|
||||
from vllm.v1.attention.backends.utils import get_kv_cache_layout
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.worker.block_table import BlockTable
|
||||
@@ -173,7 +173,7 @@ class NixlHandshakePayload(KVConnectorHandshakeMetadata):
|
||||
|
||||
|
||||
def compute_nixl_compatibility_hash(
|
||||
vllm_config: VllmConfig, attn_backend_name: str
|
||||
vllm_config: VllmConfig, attn_backend_name: str, cross_layers_blocks: bool
|
||||
) -> str:
|
||||
"""
|
||||
Compute compatibility hash for NIXL KV transfer.
|
||||
@@ -216,6 +216,7 @@ def compute_nixl_compatibility_hash(
|
||||
# Attention backend and KV cache dtype affect memory layout
|
||||
"attn_backend_name": attn_backend_name,
|
||||
"cache_dtype": str(cache_config.cache_dtype),
|
||||
"cross_layers_blocks": cross_layers_blocks,
|
||||
}
|
||||
|
||||
compat_hash = hash_factors(factors)
|
||||
@@ -298,6 +299,20 @@ class NixlConnectorMetadata(KVConnectorMetadata):
|
||||
|
||||
|
||||
class NixlConnector(KVConnectorBase_V1):
|
||||
@property
|
||||
def prefer_cross_layer_blocks(self) -> bool:
|
||||
backend = get_current_attn_backend(self._vllm_config)
|
||||
if backend.get_name() not in (
|
||||
"FLASH_ATTN",
|
||||
"FLASHINFER",
|
||||
):
|
||||
# For now there is no benefit to run cross layers when backend
|
||||
# does not support on HND
|
||||
return False
|
||||
|
||||
extra_config = self.kv_transfer_config.kv_connector_extra_config
|
||||
return bool(str(extra_config.get("enable_cross_layers_blocks", "False")))
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
@@ -309,6 +324,7 @@ class NixlConnector(KVConnectorBase_V1):
|
||||
assert vllm_config.kv_transfer_config is not None
|
||||
assert vllm_config.kv_transfer_config.engine_id is not None
|
||||
self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id
|
||||
self.kv_transfer_config = vllm_config.kv_transfer_config
|
||||
|
||||
if role == KVConnectorRole.SCHEDULER:
|
||||
self.connector_scheduler: NixlConnectorScheduler | None = (
|
||||
@@ -395,6 +411,16 @@ class NixlConnector(KVConnectorBase_V1):
|
||||
assert self.connector_worker is not None
|
||||
self.connector_worker.register_kv_caches(kv_caches)
|
||||
|
||||
def register_cross_layers_kv_cache(
|
||||
self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]
|
||||
):
|
||||
assert self.connector_worker is not None
|
||||
|
||||
cross_layer_name = "ALL_LAYERS"
|
||||
kv_caches = {cross_layer_name: kv_cache}
|
||||
|
||||
self.connector_worker.register_kv_caches(kv_caches)
|
||||
|
||||
def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
|
||||
assert self.connector_worker is not None
|
||||
self.connector_worker.set_host_xfer_buffer_ops(copy_operation)
|
||||
@@ -976,20 +1002,17 @@ class NixlConnectorWorker:
|
||||
|
||||
# Get the attention backend from the first layer
|
||||
# NOTE (NickLucche) models with multiple backends are not supported yet
|
||||
backend = get_current_attn_backend(vllm_config)
|
||||
self.attn_backend = get_current_attn_backend(vllm_config)
|
||||
|
||||
self.backend_name = backend.get_name()
|
||||
self.backend_name = self.attn_backend.get_name()
|
||||
self.kv_cache_layout = get_kv_cache_layout()
|
||||
self.host_buffer_kv_cache_layout = self.kv_cache_layout
|
||||
logger.debug("Detected attention backend %s", self.backend_name)
|
||||
logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
|
||||
|
||||
self.compat_hash = compute_nixl_compatibility_hash(
|
||||
self.vllm_config, self.backend_name
|
||||
)
|
||||
self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
|
||||
"enforce_handshake_compat", True
|
||||
)
|
||||
# lazy initialized in register_kv_caches
|
||||
self.compat_hash: str | None = None
|
||||
self.kv_topo: TpKVTopology | None = None
|
||||
|
||||
self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
|
||||
self._block_size: dict[EngineId, int] = {self.engine_id: self.block_size}
|
||||
@@ -998,17 +1021,12 @@ class NixlConnectorWorker:
|
||||
self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
|
||||
self.xfer_stats = NixlKVConnectorStats()
|
||||
|
||||
self.kv_topo = TpKVTopology(
|
||||
tp_rank=self.tp_rank,
|
||||
engine_id=self.engine_id,
|
||||
remote_tp_size=self._tp_size, # shared state
|
||||
remote_block_size=self._block_size, # shared state
|
||||
is_mla=self.use_mla,
|
||||
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
|
||||
attn_backend=backend,
|
||||
)
|
||||
self._physical_blocks_per_logical_kv_block = 1
|
||||
|
||||
self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
|
||||
"enforce_handshake_compat", True
|
||||
)
|
||||
|
||||
def _nixl_handshake(
|
||||
self,
|
||||
host: str,
|
||||
@@ -1022,6 +1040,7 @@ class NixlConnectorWorker:
|
||||
# Regardless, only handshake with the remote TP rank(s) that current
|
||||
# local rank will read from. Note that With homogeneous TP,
|
||||
# this happens to be the same single rank_i.
|
||||
assert self.kv_topo is not None
|
||||
p_remote_ranks = self.kv_topo.get_target_remote_ranks(remote_tp_size)
|
||||
remote_rank_to_agent_name = {}
|
||||
path = make_zmq_path("tcp", host, port)
|
||||
@@ -1059,6 +1078,7 @@ class NixlConnectorWorker:
|
||||
)
|
||||
|
||||
# Check compatibility hash BEFORE decoding agent metadata
|
||||
assert self.compat_hash is not None
|
||||
if (
|
||||
self.enforce_compat_hash
|
||||
and handshake_payload.compatibility_hash != self.compat_hash
|
||||
@@ -1267,6 +1287,20 @@ class NixlConnectorWorker:
|
||||
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
|
||||
"""Register the KV Cache data in nixl."""
|
||||
|
||||
self.kv_topo = TpKVTopology(
|
||||
tp_rank=self.tp_rank,
|
||||
engine_id=self.engine_id,
|
||||
remote_tp_size=self._tp_size, # shared state
|
||||
remote_block_size=self._block_size, # shared state
|
||||
is_mla=self.use_mla,
|
||||
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
|
||||
attn_backend=self.attn_backend,
|
||||
tensor_shape=next(iter(kv_caches.values())).shape,
|
||||
)
|
||||
self.compat_hash = compute_nixl_compatibility_hash(
|
||||
self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks
|
||||
)
|
||||
|
||||
if self.use_host_buffer:
|
||||
self.initialize_host_xfer_buffer(kv_caches=kv_caches)
|
||||
assert len(self.host_xfer_buffers) == len(kv_caches), (
|
||||
@@ -1301,29 +1335,21 @@ class NixlConnectorWorker:
|
||||
# (roughly 8KB vs 5KB).
|
||||
# Conversely for FlashInfer, K and V are registered in the same region
|
||||
# to better exploit the memory layout (ie num_blocks is the first dim).
|
||||
split_k_and_v = self.kv_topo.split_k_and_v
|
||||
tensor_size_bytes = None
|
||||
|
||||
# TODO (NickLucche): Get kernel_block_size in a cleaner way
|
||||
# NHD default "view" for non-MLA cache
|
||||
if self.device_type == "cpu":
|
||||
block_size_position = -2
|
||||
else:
|
||||
block_size_position = -2 if self.use_mla else -3
|
||||
|
||||
# Enable different block lengths for different layers when MLA is used.
|
||||
self.block_len_per_layer = list[int]()
|
||||
self.slot_size_per_layer = list[int]() # HD bytes in kv terms
|
||||
for layer_name, cache_or_caches in xfer_buffers.items():
|
||||
cache_list = cache_or_caches if split_k_and_v else [cache_or_caches]
|
||||
|
||||
cache_list = (
|
||||
cache_or_caches if self.kv_topo.split_k_and_v else [cache_or_caches]
|
||||
)
|
||||
for cache in cache_list:
|
||||
base_addr = cache.data_ptr()
|
||||
if base_addr in seen_base_addresses:
|
||||
continue
|
||||
|
||||
kernel_block_size = cache.shape[block_size_position]
|
||||
|
||||
kernel_block_size = cache.shape[self.kv_topo.block_size_position]
|
||||
if self.block_size != kernel_block_size:
|
||||
logger.info_once(
|
||||
"User-specified logical block size (%s) does not match"
|
||||
@@ -1385,6 +1411,7 @@ class NixlConnectorWorker:
|
||||
|
||||
self.device_kv_caches = kv_caches
|
||||
self.dst_num_blocks[self.engine_id] = self.num_blocks
|
||||
|
||||
if self.kv_topo.is_kv_layout_blocks_first:
|
||||
for i in range(len(self.slot_size_per_layer)):
|
||||
assert self.slot_size_per_layer[i] % 2 == 0
|
||||
@@ -1440,6 +1467,7 @@ class NixlConnectorWorker:
|
||||
block_size=self.block_size,
|
||||
)
|
||||
# Wrap metadata in payload with hash for defensive decoding
|
||||
assert self.compat_hash is not None
|
||||
encoder = msgspec.msgpack.Encoder()
|
||||
self.xfer_handshake_metadata = NixlHandshakePayload(
|
||||
compatibility_hash=self.compat_hash,
|
||||
@@ -1461,6 +1489,8 @@ class NixlConnectorWorker:
|
||||
register another local_xfer_handler using remote block len to ensure
|
||||
data copy correctness.
|
||||
"""
|
||||
assert self.kv_topo is not None
|
||||
|
||||
block_size_ratio = self.block_size // block_size
|
||||
blocks_data = []
|
||||
for i, base_addr in enumerate(self.seen_base_addresses):
|
||||
@@ -1573,6 +1603,7 @@ class NixlConnectorWorker:
|
||||
# remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|
|
||||
# local origin:| 0| 1| 8| 12|
|
||||
# local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|
|
||||
assert self.kv_topo is not None
|
||||
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(engine_id)
|
||||
|
||||
if engine_id not in self.dst_num_blocks:
|
||||
@@ -1700,7 +1731,10 @@ class NixlConnectorWorker:
|
||||
"""
|
||||
remote_engine_id = nixl_agent_meta.engine_id
|
||||
|
||||
assert self._tp_size[remote_engine_id] == remote_tp_size
|
||||
assert (
|
||||
self._tp_size[remote_engine_id] == remote_tp_size
|
||||
and self.kv_topo is not None
|
||||
)
|
||||
|
||||
tp_ratio = self.kv_topo.tp_ratio_from_engine_id(remote_engine_id)
|
||||
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
|
||||
@@ -1837,6 +1871,7 @@ class NixlConnectorWorker:
|
||||
if len(self.device_kv_caches) == 0:
|
||||
return
|
||||
assert block_size_ratio >= 1, "Only nP < nD supported currently."
|
||||
assert self.kv_topo is not None
|
||||
if self.enable_permute_local_kv and block_size_ratio > 1:
|
||||
logger.debug(
|
||||
"Post-processing device kv cache on receive by converting "
|
||||
@@ -1856,7 +1891,7 @@ class NixlConnectorWorker:
|
||||
block_size_ratio,
|
||||
)
|
||||
|
||||
split_k_and_v = not (self.use_mla or self.kv_topo.is_kv_layout_blocks_first)
|
||||
split_k_and_v = self.kv_topo.split_k_and_v
|
||||
|
||||
for block_ids in block_ids_list:
|
||||
indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
|
||||
@@ -1881,6 +1916,7 @@ class NixlConnectorWorker:
|
||||
The scheduler process (via the MultiprocExecutor) will use this output
|
||||
to track which workers are done.
|
||||
"""
|
||||
assert self.kv_topo is not None
|
||||
done_sending = self._get_new_notifs()
|
||||
done_recving = self._pop_done_transfers(self._recving_transfers)
|
||||
|
||||
@@ -1950,6 +1986,7 @@ class NixlConnectorWorker:
|
||||
are reading from the same producer (heterogeneous TP scenario), wait
|
||||
for all consumers to be done pulling.
|
||||
"""
|
||||
assert self.kv_topo is not None
|
||||
notified_req_ids: set[str] = set()
|
||||
for notifs in self.nixl_wrapper.get_new_notifs().values():
|
||||
for notif in notifs:
|
||||
@@ -2109,7 +2146,7 @@ class NixlConnectorWorker:
|
||||
self._reqs_to_send[req_id] = expiration_time
|
||||
|
||||
def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
|
||||
assert meta.remote is not None
|
||||
assert meta.remote is not None and self.kv_topo is not None
|
||||
remote_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(
|
||||
meta.remote.engine_id
|
||||
)
|
||||
@@ -2178,10 +2215,7 @@ class NixlConnectorWorker:
|
||||
local_xfer_side_handle: int,
|
||||
remote_xfer_side_handle: int,
|
||||
):
|
||||
"""
|
||||
Post a READ point-to-point xfer request from a single local worker to
|
||||
a single remote worker.
|
||||
"""
|
||||
assert self.kv_topo is not None
|
||||
block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(dst_engine_id)
|
||||
if block_size_ratio > 1:
|
||||
local_block_ids = self.get_mapped_blocks(
|
||||
@@ -2414,6 +2448,7 @@ class NixlConnectorWorker:
|
||||
For FlashInfer, this is half the length of the whole block, as K and V
|
||||
share the same region.
|
||||
"""
|
||||
assert self.kv_topo is not None
|
||||
if self.kv_topo.is_kv_layout_blocks_first:
|
||||
# For indexing only half (either just the K or V part).
|
||||
block_len = self.block_len_per_layer[layer_idx] // 2
|
||||
|
||||
@@ -1000,7 +1000,7 @@ class GroupCoordinator:
|
||||
if self.device_communicator is not None:
|
||||
self.device_communicator.prepare_communication_buffer_for_model(model)
|
||||
|
||||
def dispatch(
|
||||
def dispatch_router_logits(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
@@ -1011,7 +1011,7 @@ class GroupCoordinator:
|
||||
| tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
):
|
||||
if self.device_communicator is not None:
|
||||
return self.device_communicator.dispatch( # type: ignore[call-arg]
|
||||
return self.device_communicator.dispatch_router_logits(
|
||||
hidden_states,
|
||||
router_logits,
|
||||
is_sequence_parallel,
|
||||
@@ -1020,6 +1020,28 @@ class GroupCoordinator:
|
||||
else:
|
||||
return hidden_states, router_logits
|
||||
|
||||
def dispatch(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
is_sequence_parallel: bool = False,
|
||||
extra_tensors: list[torch.Tensor] | None = None,
|
||||
) -> (
|
||||
tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
|
||||
| tuple[torch.Tensor, torch.Tensor, torch.Tensor]
|
||||
):
|
||||
if self.device_communicator is not None:
|
||||
return self.device_communicator.dispatch(
|
||||
hidden_states,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
is_sequence_parallel,
|
||||
extra_tensors,
|
||||
)
|
||||
else:
|
||||
return hidden_states, topk_weights, topk_ids
|
||||
|
||||
def combine(
|
||||
self, hidden_states, is_sequence_parallel: bool = False
|
||||
) -> torch.Tensor:
|
||||
|
||||
@@ -264,6 +264,39 @@ def load_log_config(log_config_file: str | None) -> dict | None:
|
||||
return None
|
||||
|
||||
|
||||
def get_uvicorn_log_config(args: Namespace) -> dict | None:
|
||||
"""
|
||||
Get the uvicorn log config based on the provided arguments.
|
||||
|
||||
Priority:
|
||||
1. If log_config_file is specified, use it
|
||||
2. If disable_access_log_for_endpoints is specified, create a config with
|
||||
the access log filter
|
||||
3. Otherwise, return None (use uvicorn defaults)
|
||||
"""
|
||||
# First, try to load from file if specified
|
||||
log_config = load_log_config(args.log_config_file)
|
||||
if log_config is not None:
|
||||
return log_config
|
||||
|
||||
# If endpoints to filter are specified, create a config with the filter
|
||||
if args.disable_access_log_for_endpoints:
|
||||
from vllm.logging_utils import create_uvicorn_log_config
|
||||
|
||||
# Parse comma-separated string into list
|
||||
excluded_paths = [
|
||||
p.strip()
|
||||
for p in args.disable_access_log_for_endpoints.split(",")
|
||||
if p.strip()
|
||||
]
|
||||
return create_uvicorn_log_config(
|
||||
excluded_paths=excluded_paths,
|
||||
log_level=args.uvicorn_log_level,
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class AuthenticationMiddleware:
|
||||
"""
|
||||
Pure ASGI middleware that authenticates each request by checking
|
||||
@@ -930,8 +963,8 @@ async def run_server_worker(
|
||||
if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3:
|
||||
ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin)
|
||||
|
||||
# Load logging config for uvicorn if specified
|
||||
log_config = load_log_config(args.log_config_file)
|
||||
# Get uvicorn log config (from file or with endpoint filter)
|
||||
log_config = get_uvicorn_log_config(args)
|
||||
if log_config is not None:
|
||||
uvicorn_kwargs["log_config"] = log_config
|
||||
|
||||
|
||||
@@ -44,6 +44,7 @@ from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaMessage,
|
||||
DeltaToolCall,
|
||||
ErrorResponse,
|
||||
FunctionCall,
|
||||
PromptTokenUsageInfo,
|
||||
RequestResponseMetadata,
|
||||
ToolCall,
|
||||
@@ -143,11 +144,6 @@ class OpenAIServingChat(OpenAIServing):
|
||||
self.enable_prompt_tokens_details = enable_prompt_tokens_details
|
||||
self.enable_force_include_usage = enable_force_include_usage
|
||||
self.default_sampling_params = self.model_config.get_diff_sampling_param()
|
||||
if self.model_config.hf_config.model_type == "kimi_k2":
|
||||
self.tool_call_id_type = "kimi_k2"
|
||||
else:
|
||||
self.tool_call_id_type = "random"
|
||||
|
||||
self.use_harmony = self.model_config.hf_config.model_type == "gpt_oss"
|
||||
if self.use_harmony:
|
||||
if "stop_token_ids" not in self.default_sampling_params:
|
||||
@@ -156,6 +152,16 @@ class OpenAIServingChat(OpenAIServing):
|
||||
get_stop_tokens_for_assistant_actions()
|
||||
)
|
||||
|
||||
# Handle tool call ID type for Kimi K2 (supporting test mocking via overrides)
|
||||
hf_overrides = getattr(self.model_config, "hf_overrides", None)
|
||||
if self.model_config.hf_text_config.model_type == "kimi_k2" or (
|
||||
isinstance(hf_overrides, dict)
|
||||
and hf_overrides.get("model_type") == "kimi_k2"
|
||||
):
|
||||
self.tool_call_id_type = "kimi_k2"
|
||||
else:
|
||||
self.tool_call_id_type = "random"
|
||||
|
||||
# NOTE(woosuk): While OpenAI's chat completion API supports browsing
|
||||
# for some models, currently vLLM doesn't support it. Please use the
|
||||
# Responses API instead.
|
||||
@@ -247,8 +253,8 @@ class OpenAIServingChat(OpenAIServing):
|
||||
# because of issues with pydantic we need to potentially
|
||||
# re-serialize the tool_calls field of the request
|
||||
# for more info: see comment in `maybe_serialize_tool_calls`
|
||||
maybe_serialize_tool_calls(request)
|
||||
truncate_tool_call_ids(request)
|
||||
maybe_serialize_tool_calls(request) # type: ignore[arg-type]
|
||||
truncate_tool_call_ids(request) # type: ignore[arg-type]
|
||||
validate_request_params(request)
|
||||
|
||||
# Check if tool parsing is unavailable (common condition)
|
||||
@@ -454,6 +460,7 @@ class OpenAIServingChat(OpenAIServing):
|
||||
|
||||
# Streaming response
|
||||
tokenizer = self.renderer.tokenizer
|
||||
assert tokenizer is not None
|
||||
|
||||
if request.stream:
|
||||
return self.chat_completion_stream_generator(
|
||||
@@ -632,9 +639,11 @@ class OpenAIServingChat(OpenAIServing):
|
||||
request_id: str,
|
||||
model_name: str,
|
||||
conversation: list[ConversationMessage],
|
||||
tokenizer: TokenizerLike | None,
|
||||
tokenizer: TokenizerLike,
|
||||
request_metadata: RequestResponseMetadata,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
from vllm.tokenizers.mistral import MistralTokenizer
|
||||
|
||||
created_time = int(time.time())
|
||||
chunk_object_type: Final = "chat.completion.chunk"
|
||||
first_iteration = True
|
||||
@@ -698,7 +707,7 @@ class OpenAIServingChat(OpenAIServing):
|
||||
)
|
||||
reasoning_parser = self.reasoning_parser(
|
||||
tokenizer,
|
||||
chat_template_kwargs=chat_template_kwargs, # type: ignore[call-arg]
|
||||
chat_template_kwargs=chat_template_kwargs or {}, # type: ignore[call-arg]
|
||||
)
|
||||
except RuntimeError as e:
|
||||
logger.exception("Error in reasoning parser creation.")
|
||||
@@ -955,8 +964,17 @@ class OpenAIServingChat(OpenAIServing):
|
||||
index=i,
|
||||
)
|
||||
else:
|
||||
# Generate ID based on tokenizer type
|
||||
if isinstance(tokenizer, MistralTokenizer):
|
||||
tool_call_id = MistralToolCall.generate_random_id()
|
||||
else:
|
||||
tool_call_id = make_tool_call_id(
|
||||
id_type=self.tool_call_id_type,
|
||||
func_name=tool_choice_function_name,
|
||||
idx=history_tool_call_cnt,
|
||||
)
|
||||
delta_tool_call = DeltaToolCall(
|
||||
id=make_tool_call_id(),
|
||||
id=tool_call_id,
|
||||
type="function",
|
||||
function=DeltaFunctionCall(
|
||||
name=tool_choice_function_name,
|
||||
@@ -1387,9 +1405,11 @@ class OpenAIServingChat(OpenAIServing):
|
||||
request_id: str,
|
||||
model_name: str,
|
||||
conversation: list[ConversationMessage],
|
||||
tokenizer: TokenizerLike | None,
|
||||
tokenizer: TokenizerLike,
|
||||
request_metadata: RequestResponseMetadata,
|
||||
) -> ErrorResponse | ChatCompletionResponse:
|
||||
from vllm.tokenizers.mistral import MistralTokenizer
|
||||
|
||||
created_time = int(time.time())
|
||||
final_res: RequestOutput | None = None
|
||||
|
||||
@@ -1524,39 +1544,85 @@ class OpenAIServingChat(OpenAIServing):
|
||||
tool_call_class = (
|
||||
MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall
|
||||
)
|
||||
if (not self.enable_auto_tools or not self.tool_parser) and (
|
||||
if self.use_harmony:
|
||||
# Harmony models already have parsed content and tool_calls
|
||||
# through parse_chat_output. Respect its output directly.
|
||||
message = ChatMessage(
|
||||
role=role,
|
||||
reasoning=reasoning,
|
||||
content=content,
|
||||
tool_calls=tool_calls if tool_calls else [],
|
||||
)
|
||||
|
||||
elif (not self.enable_auto_tools or not self.tool_parser) and (
|
||||
not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
|
||||
and request.tool_choice != "required"
|
||||
):
|
||||
message = ChatMessage(role=role, reasoning=reasoning, content=content)
|
||||
|
||||
# if the request uses tools and specified a tool choice
|
||||
elif (
|
||||
request.tool_choice
|
||||
and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam
|
||||
):
|
||||
assert tool_calls is not None and len(tool_calls) > 0
|
||||
tool_call_class_items = []
|
||||
for idx, tc in enumerate(tool_calls):
|
||||
# Use native ID if available (e.g., Kimi K2),
|
||||
# otherwise generate ID with correct id_type
|
||||
if tc.id:
|
||||
tool_call_class_items.append(
|
||||
tool_call_class(id=tc.id, function=tc)
|
||||
)
|
||||
else:
|
||||
# Generate ID using the correct format (kimi_k2 or random),
|
||||
# but leave it to the class if it's Mistral to preserve
|
||||
# 9-char IDs
|
||||
if isinstance(tokenizer, MistralTokenizer):
|
||||
tool_call_class_items.append(tool_call_class(function=tc))
|
||||
else:
|
||||
generated_id = make_tool_call_id(
|
||||
id_type=self.tool_call_id_type,
|
||||
func_name=tc.name,
|
||||
idx=history_tool_call_cnt + idx,
|
||||
)
|
||||
tool_call_class_items.append(
|
||||
tool_call_class(id=generated_id, function=tc)
|
||||
)
|
||||
history_tool_call_cnt += 1
|
||||
message = ChatMessage(
|
||||
role=role,
|
||||
reasoning=reasoning,
|
||||
content="",
|
||||
tool_calls=[tool_call_class(function=tc) for tc in tool_calls],
|
||||
tool_calls=tool_call_class_items,
|
||||
)
|
||||
|
||||
elif request.tool_choice and request.tool_choice == "required":
|
||||
tool_call_class_items = []
|
||||
assert tool_calls is not None and len(tool_calls) > 0
|
||||
for tool_call in tool_calls:
|
||||
tool_call_class_items.append(
|
||||
tool_call_class(
|
||||
id=make_tool_call_id(
|
||||
for idx, tool_call in enumerate(tool_calls):
|
||||
# Use native ID if available,
|
||||
# otherwise generate ID with correct id_type
|
||||
if tool_call.id:
|
||||
tool_call_class_items.append(
|
||||
tool_call_class(id=tool_call.id, function=tool_call)
|
||||
)
|
||||
else:
|
||||
# Generate ID using the correct format (kimi_k2 or random),
|
||||
# but leave it to the class if it's Mistral to preserve
|
||||
# 9-char IDs
|
||||
if isinstance(tokenizer, MistralTokenizer):
|
||||
tool_call_class_items.append(
|
||||
tool_call_class(function=tool_call)
|
||||
)
|
||||
else:
|
||||
generated_id = make_tool_call_id(
|
||||
id_type=self.tool_call_id_type,
|
||||
func_name=tool_call.name,
|
||||
idx=history_tool_call_cnt,
|
||||
),
|
||||
function=tool_call,
|
||||
)
|
||||
)
|
||||
idx=history_tool_call_cnt + idx,
|
||||
)
|
||||
tool_call_class_items.append(
|
||||
tool_call_class(id=generated_id, function=tool_call)
|
||||
)
|
||||
history_tool_call_cnt += 1
|
||||
message = ChatMessage(
|
||||
role=role,
|
||||
@@ -1582,17 +1648,35 @@ class OpenAIServingChat(OpenAIServing):
|
||||
# call. The same is not true for named function calls
|
||||
auto_tools_called = tool_calls is not None and len(tool_calls) > 0
|
||||
if tool_calls:
|
||||
tool_call_items = []
|
||||
for idx, tc in enumerate(tool_calls):
|
||||
# Use native ID if available (e.g., Kimi K2),
|
||||
# otherwise generate ID with correct id_type
|
||||
if tc.id:
|
||||
tool_call_items.append(
|
||||
tool_call_class(id=tc.id, function=tc)
|
||||
)
|
||||
else:
|
||||
# Generate ID using the correct format (kimi_k2 or random),
|
||||
# but leave it to the class if it's Mistral to preserve
|
||||
# 9-char IDs
|
||||
if isinstance(tokenizer, MistralTokenizer):
|
||||
tool_call_items.append(tool_call_class(function=tc))
|
||||
else:
|
||||
generated_id = make_tool_call_id(
|
||||
id_type=self.tool_call_id_type,
|
||||
func_name=tc.name,
|
||||
idx=history_tool_call_cnt + idx,
|
||||
)
|
||||
tool_call_items.append(
|
||||
tool_call_class(id=generated_id, function=tc)
|
||||
)
|
||||
history_tool_call_cnt += 1
|
||||
message = ChatMessage(
|
||||
role=role,
|
||||
reasoning=reasoning,
|
||||
content=content,
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
function=tc,
|
||||
type="function",
|
||||
)
|
||||
for tc in tool_calls
|
||||
],
|
||||
tool_calls=tool_call_items,
|
||||
)
|
||||
|
||||
else:
|
||||
@@ -1701,13 +1785,11 @@ class OpenAIServingChat(OpenAIServing):
|
||||
elif choice.message.tool_calls:
|
||||
# For tool calls, log the function name and arguments
|
||||
tool_call_descriptions = []
|
||||
for tc in choice.message.tool_calls:
|
||||
if hasattr(tc.function, "name") and hasattr(
|
||||
tc.function, "arguments"
|
||||
):
|
||||
tool_call_descriptions.append(
|
||||
f"{tc.function.name}({tc.function.arguments})"
|
||||
)
|
||||
for tc in choice.message.tool_calls: # type: ignore
|
||||
function_call: FunctionCall = tc.function # type: ignore
|
||||
tool_call_descriptions.append(
|
||||
f"{function_call.name}({function_call.arguments})"
|
||||
)
|
||||
tool_calls_str = ", ".join(tool_call_descriptions)
|
||||
output_text = f"[tool_calls: {tool_calls_str}]"
|
||||
|
||||
@@ -1895,7 +1977,7 @@ class OpenAIServingChat(OpenAIServing):
|
||||
# because of issues with pydantic we need to potentially
|
||||
# re-serialize the tool_calls field of the request
|
||||
# for more info: see comment in `maybe_serialize_tool_calls`
|
||||
maybe_serialize_tool_calls(request)
|
||||
maybe_serialize_tool_calls(request) # type: ignore[arg-type]
|
||||
|
||||
# Add system message.
|
||||
# NOTE: In Chat Completion API, browsing is enabled by default
|
||||
@@ -1913,7 +1995,7 @@ class OpenAIServingChat(OpenAIServing):
|
||||
# Add developer message.
|
||||
if request.tools:
|
||||
dev_msg = get_developer_message(
|
||||
tools=request.tools if should_include_tools else None
|
||||
tools=request.tools if should_include_tools else None # type: ignore[arg-type]
|
||||
)
|
||||
messages.append(dev_msg)
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user