forked from Karylab-cklius/vllm
Compare commits
121
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9df42d2000 | ||
|
|
8d151aa148 | ||
|
|
11a13b0fd3 | ||
|
|
df3f537a66 | ||
|
|
7743152957 | ||
|
|
ab33d2a629 | ||
|
|
be3af2d29e | ||
|
|
c656ba3b4d | ||
|
|
dc5fa77a4e | ||
|
|
1e4a084c8e | ||
|
|
7967e854da | ||
|
|
6bd6d0c3c1 | ||
|
|
8e962fef5f | ||
|
|
574fe75245 | ||
|
|
c61a98f529 | ||
|
|
28bffe9466 | ||
|
|
ad65177a19 | ||
|
|
d44a5b6c47 | ||
|
|
1d65283e95 | ||
|
|
c464b57374 | ||
|
|
c5c38e152a | ||
|
|
d00df624f3 | ||
|
|
9752da9d9c | ||
|
|
04925b2202 | ||
|
|
d74278fb67 | ||
|
|
b68fd899d1 | ||
|
|
0b5f9b7204 | ||
|
|
9a8853f781 | ||
|
|
387a1898d9 | ||
|
|
3b30e61507 | ||
|
|
824f9e8f3c | ||
|
|
6cc403e67d | ||
|
|
72d5951d02 | ||
|
|
a3205beffb | ||
|
|
6930becd45 | ||
|
|
03a8770a6d | ||
|
|
bc56a1d56e | ||
|
|
ec7d9e6745 | ||
|
|
3bb4e4311c | ||
|
|
08f8c198ae | ||
|
|
a21cedf4ff | ||
|
|
3ef74cde5d | ||
|
|
cd81cdb399 | ||
|
|
1e828573b4 | ||
|
|
a5ccc85c8c | ||
|
|
b5475d0534 | ||
|
|
9521002f0a | ||
|
|
ec17bdd894 | ||
|
|
bb59c90248 | ||
|
|
5bff999d12 | ||
|
|
bb85929aa6 | ||
|
|
5653021094 | ||
|
|
974d829b05 | ||
|
|
91ac5d9bfd | ||
|
|
23d825aba1 | ||
|
|
f07a128413 | ||
|
|
71cd89264f | ||
|
|
19fab44152 | ||
|
|
79c7e09235 | ||
|
|
79f3fab05a | ||
|
|
604b9eaec5 | ||
|
|
50dbd6c9e6 | ||
|
|
98bcc6ca59 | ||
|
|
f13e86d8dd | ||
|
|
9ca768c740 | ||
|
|
d5fe3f702c | ||
|
|
73391a1baa | ||
|
|
b3c14229b0 | ||
|
|
2f186635cb | ||
|
|
342a7cda2d | ||
|
|
d1ea65d0a1 | ||
|
|
de42abb366 | ||
|
|
60ca7981bc | ||
|
|
0ef5b9147b | ||
|
|
ed242652d7 | ||
|
|
b37b679770 | ||
|
|
a0638d052d | ||
|
|
c027541eaf | ||
|
|
fd267bc7b7 | ||
|
|
bfaa559305 | ||
|
|
87789c8364 | ||
|
|
bcd65c1f6a | ||
|
|
59d53066d8 | ||
|
|
4a9952ec1b | ||
|
|
1dae7b7843 | ||
|
|
5885e330ef | ||
|
|
071d863e20 | ||
|
|
0916e7960b | ||
|
|
3d2a026fd0 | ||
|
|
dddbff4624 | ||
|
|
47e9b63e1a | ||
|
|
934acddef9 | ||
|
|
742d214d6e | ||
|
|
4137c5dfa7 | ||
|
|
7a8a46ddcb | ||
|
|
bcf0731aa0 | ||
|
|
ec090c2429 | ||
|
|
eea3024f43 | ||
|
|
2f308214c0 | ||
|
|
1b4e8e53f8 | ||
|
|
dcf6ee8592 | ||
|
|
372b2e762a | ||
|
|
6afa587d31 | ||
|
|
94ed6cf6ea | ||
|
|
bf37812ca7 | ||
|
|
b86bf4417e | ||
|
|
de13dd781f | ||
|
|
62788f99a4 | ||
|
|
ea5ff3a1f6 | ||
|
|
04ea31baab | ||
|
|
6f019e6e0a | ||
|
|
d707678dfb | ||
|
|
fc22cae4ac | ||
|
|
96161fe978 | ||
|
|
4453ba8d9e | ||
|
|
aa181c923b | ||
|
|
be7370daf3 | ||
|
|
9ea1f598ce | ||
|
|
f120bd42d3 | ||
|
|
fac4e96940 | ||
|
|
6d4e27ce29 |
@@ -8,7 +8,7 @@ clean_docker_tag() {
|
||||
}
|
||||
|
||||
print_usage_and_exit() {
|
||||
echo "Usage: $0 <registry> <repo> <commit> <branch> <vllm_use_precompiled> <vllm_merge_base_commit> <cache_from> <cache_to>"
|
||||
echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag> [<image_tag_latest>]"
|
||||
exit 1
|
||||
}
|
||||
|
||||
@@ -142,11 +142,16 @@ resolve_parent_commit() {
|
||||
|
||||
print_bake_config() {
|
||||
echo "--- :page_facing_up: Resolved bake configuration"
|
||||
BAKE_CONFIG_FILE="bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
|
||||
# Write to a temp directory to avoid polluting the repo root (which is the
|
||||
# Docker build context). Files left in the repo root get COPY'd into the
|
||||
# image and can cause duplicate artifact uploads from downstream steps.
|
||||
local bake_tmp
|
||||
bake_tmp="$(mktemp -d)"
|
||||
BAKE_CONFIG_FILE="${bake_tmp}/bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
|
||||
docker buildx bake -f "${VLLM_BAKE_FILE_PATH}" -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}"
|
||||
(cd "$(dirname "${BAKE_CONFIG_FILE}")" && buildkite-agent artifact upload "$(basename "${BAKE_CONFIG_FILE}")")
|
||||
}
|
||||
|
||||
#################################
|
||||
@@ -154,7 +159,7 @@ print_bake_config() {
|
||||
#################################
|
||||
print_instance_info
|
||||
|
||||
if [[ $# -lt 7 ]]; then
|
||||
if [[ $# -lt 5 ]]; then
|
||||
print_usage_and_exit
|
||||
fi
|
||||
|
||||
@@ -163,10 +168,8 @@ REGISTRY=$1
|
||||
REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
BRANCH=$4
|
||||
VLLM_USE_PRECOMPILED=$5
|
||||
VLLM_MERGE_BASE_COMMIT=$6
|
||||
IMAGE_TAG=$7
|
||||
IMAGE_TAG_LATEST=${8:-} # only used for main branch, optional
|
||||
IMAGE_TAG=$5
|
||||
IMAGE_TAG_LATEST=${6:-} # only used for main branch, optional
|
||||
|
||||
# build config
|
||||
TARGET="test-ci"
|
||||
@@ -193,8 +196,6 @@ 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"
|
||||
@@ -202,8 +203,6 @@ 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}"
|
||||
|
||||
|
||||
@@ -5,8 +5,7 @@ steps:
|
||||
depends_on: []
|
||||
timeout_in_minutes: 600
|
||||
commands:
|
||||
- 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 $IMAGE_TAG_LATEST; fi
|
||||
- if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG $IMAGE_TAG_LATEST; else .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG; fi
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: -1 # Agent was lost
|
||||
|
||||
@@ -11,10 +11,10 @@ REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
|
||||
# authenticate with AWS ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
|
||||
|
||||
# skip build if image already exists
|
||||
if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu) ]]; then
|
||||
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then
|
||||
echo "Image not found, proceeding with build..."
|
||||
else
|
||||
echo "Image found"
|
||||
@@ -24,13 +24,13 @@ fi
|
||||
# build
|
||||
docker build --file docker/Dockerfile.cpu \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg buildkite_commit=$BUILDKITE_COMMIT \
|
||||
--build-arg buildkite_commit="$BUILDKITE_COMMIT" \
|
||||
--build-arg VLLM_CPU_AVX512BF16=true \
|
||||
--build-arg VLLM_CPU_AVX512VNNI=true \
|
||||
--build-arg VLLM_CPU_AMXBF16=true \
|
||||
--tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu \
|
||||
--tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu \
|
||||
--target vllm-test \
|
||||
--progress plain .
|
||||
|
||||
# push
|
||||
docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu
|
||||
docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu
|
||||
|
||||
@@ -11,10 +11,10 @@ REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
|
||||
# authenticate with AWS ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
|
||||
|
||||
# skip build if image already exists
|
||||
if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu) ]]; then
|
||||
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then
|
||||
echo "Image not found, proceeding with build..."
|
||||
else
|
||||
echo "Image found"
|
||||
@@ -24,10 +24,10 @@ fi
|
||||
# build
|
||||
docker build --file docker/Dockerfile.cpu \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg buildkite_commit=$BUILDKITE_COMMIT \
|
||||
--tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu \
|
||||
--build-arg buildkite_commit="$BUILDKITE_COMMIT" \
|
||||
--tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu \
|
||||
--target vllm-test \
|
||||
--progress plain .
|
||||
|
||||
# push
|
||||
docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu
|
||||
docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu
|
||||
|
||||
@@ -11,10 +11,10 @@ REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
|
||||
# authenticate with AWS ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
|
||||
|
||||
# skip build if image already exists
|
||||
if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu) ]]; then
|
||||
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu) ]]; then
|
||||
echo "Image not found, proceeding with build..."
|
||||
else
|
||||
echo "Image found"
|
||||
@@ -25,10 +25,10 @@ fi
|
||||
docker build \
|
||||
--file tests/pytorch_ci_hud_benchmark/Dockerfile.hpu \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg buildkite_commit=$BUILDKITE_COMMIT \
|
||||
--tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu \
|
||||
--build-arg buildkite_commit="$BUILDKITE_COMMIT" \
|
||||
--tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu \
|
||||
--progress plain \
|
||||
https://github.com/vllm-project/vllm-gaudi.git
|
||||
|
||||
# push
|
||||
docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu
|
||||
docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# We can use this script to compute baseline accuracy on chartqa for vllm.
|
||||
#
|
||||
# Make sure you have lm-eval-harness installed:
|
||||
# pip install "lm-eval[api]>=0.4.9.2"
|
||||
# pip install "lm-eval[api]>=0.4.11"
|
||||
|
||||
usage() {
|
||||
echo``
|
||||
@@ -41,4 +41,4 @@ lm_eval --model vllm-vlm \
|
||||
--tasks chartqa \
|
||||
--batch_size auto \
|
||||
--apply_chat_template \
|
||||
--limit $LIMIT
|
||||
--limit "$LIMIT"
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# We can use this script to compute baseline accuracy on GSM for transformers.
|
||||
#
|
||||
# Make sure you have lm-eval-harness installed:
|
||||
# pip install "lm-eval[api]>=0.4.9.2"
|
||||
# pip install "lm-eval[api]>=0.4.11"
|
||||
|
||||
usage() {
|
||||
echo``
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
# We use this for fp8, which HF does not support.
|
||||
#
|
||||
# Make sure you have lm-eval-harness installed:
|
||||
# pip install "lm-eval[api]>=0.4.9.2"
|
||||
# pip install "lm-eval[api]>=0.4.11"
|
||||
|
||||
usage() {
|
||||
echo``
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
# We use this for fp8, which HF does not support.
|
||||
#
|
||||
# Make sure you have lm-eval-harness installed:
|
||||
# pip install "lm-eval[api]>=0.4.9.2"
|
||||
# pip install "lm-eval[api]>=0.4.11"
|
||||
|
||||
usage() {
|
||||
echo``
|
||||
@@ -20,14 +20,11 @@ usage() {
|
||||
echo
|
||||
}
|
||||
|
||||
while getopts "m:b:l:f:t:" OPT; do
|
||||
while getopts "m:l:f:t:" OPT; do
|
||||
case ${OPT} in
|
||||
m )
|
||||
MODEL="$OPTARG"
|
||||
;;
|
||||
b )
|
||||
BATCH_SIZE="$OPTARG"
|
||||
;;
|
||||
l )
|
||||
LIMIT="$OPTARG"
|
||||
;;
|
||||
|
||||
@@ -15,11 +15,11 @@ DTYPE_FILTER="${DTYPE_FILTER:-}"
|
||||
check_gpus() {
|
||||
if command -v nvidia-smi; then
|
||||
# check the number of GPUs and GPU type.
|
||||
declare -g gpu_count=$(nvidia-smi --list-gpus | wc -l)
|
||||
declare -g gpu_count=$(nvidia-smi --list-gpus | grep -c . || true)
|
||||
elif command -v amd-smi; then
|
||||
declare -g gpu_count=$(amd-smi list | grep 'GPU' | wc -l)
|
||||
declare -g gpu_count=$(amd-smi list | grep -c 'GPU' || true)
|
||||
elif command -v hl-smi; then
|
||||
declare -g gpu_count=$(hl-smi --list | grep -i "Module ID" | wc -l)
|
||||
declare -g gpu_count=$(hl-smi --list | grep -ci "Module ID" || true)
|
||||
fi
|
||||
|
||||
if [[ $gpu_count -gt 0 ]]; then
|
||||
@@ -47,7 +47,7 @@ check_cpus() {
|
||||
declare -g numa_count=$(lscpu | grep "NUMA node(s):" | awk '{print $3}')
|
||||
if [[ $numa_count -gt 0 ]]; then
|
||||
echo "NUMA found."
|
||||
echo $numa_count
|
||||
echo "$numa_count"
|
||||
else
|
||||
echo "Need at least 1 NUMA to run benchmarking."
|
||||
exit 1
|
||||
@@ -434,7 +434,7 @@ run_serving_tests() {
|
||||
|
||||
# iterate over different max_concurrency
|
||||
for max_concurrency in $max_concurrency_list; do
|
||||
new_test_name=$test_name"_qps_"$qps"_concurrency_"$max_concurrency
|
||||
new_test_name="${test_name}_qps_${qps}_concurrency_${max_concurrency}"
|
||||
echo " new test name $new_test_name"
|
||||
# pass the tensor parallel size, the compilation mode, and the optimization
|
||||
# level to the client so that they can be used on the benchmark dashboard
|
||||
@@ -471,7 +471,7 @@ run_serving_tests() {
|
||||
|
||||
# clean up
|
||||
if [[ "${DRY_RUN:-0}" != "1" ]]; then
|
||||
kill -9 $server_pid
|
||||
kill -9 "$server_pid"
|
||||
kill_gpu_processes
|
||||
fi
|
||||
done
|
||||
|
||||
@@ -31,7 +31,7 @@ steps:
|
||||
commands:
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
|
||||
@@ -70,7 +70,7 @@ steps:
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
|
||||
@@ -123,7 +123,7 @@ steps:
|
||||
queue: cpu_queue_postmerge
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
|
||||
# re-tag to default image tag and push, just in case arm64 build fails
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
|
||||
@@ -137,7 +137,7 @@ steps:
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
# compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILDER_CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
|
||||
|
||||
- block: "Build release image for x86_64 CPU"
|
||||
|
||||
@@ -25,7 +25,7 @@ S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}"
|
||||
S3_URL="http://${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com"
|
||||
|
||||
# Format ROCm version for path (e.g., "7.1" -> "rocm710")
|
||||
ROCM_VERSION_PATH="rocm$(echo ${ROCM_VERSION} | tr -d '.')"
|
||||
ROCM_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
|
||||
|
||||
@@ -83,7 +83,7 @@ case "${1:-}" in
|
||||
exit 1
|
||||
fi
|
||||
|
||||
WHEEL_COUNT=$(ls artifacts/rocm-base-wheels/*.whl 2>/dev/null | wc -l)
|
||||
WHEEL_COUNT=$(find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
|
||||
if [[ "$WHEEL_COUNT" -eq 0 ]]; then
|
||||
echo "ERROR: No wheels found in artifacts/rocm-base-wheels/" >&2
|
||||
exit 1
|
||||
@@ -110,9 +110,9 @@ case "${1:-}" in
|
||||
|
||||
echo ""
|
||||
echo "Downloaded wheels:"
|
||||
ls -lh artifacts/rocm-base-wheels/
|
||||
find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' -exec ls -lh {} \;
|
||||
|
||||
WHEEL_COUNT=$(ls artifacts/rocm-base-wheels/*.whl 2>/dev/null | wc -l)
|
||||
WHEEL_COUNT=$(find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
|
||||
echo ""
|
||||
echo "Total: $WHEEL_COUNT wheels"
|
||||
echo "========================================"
|
||||
|
||||
@@ -134,7 +134,7 @@ log_info "Fetching merged PRs from milestone '${MILESTONE}'..."
|
||||
|
||||
# Store PR data in a temp file
|
||||
PR_DATA=$(mktemp)
|
||||
trap "rm -f $PR_DATA" EXIT
|
||||
trap 'rm -f "$PR_DATA"' EXIT
|
||||
|
||||
if ! gh pr list --state merged --search "milestone:${MILESTONE}" \
|
||||
--limit 1000 \
|
||||
|
||||
@@ -27,7 +27,7 @@ function cpu_tests() {
|
||||
podman exec -it "$container_id" bash -c "
|
||||
export TORCH_COMPILE_DISABLE=1
|
||||
set -xve
|
||||
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m" >> $HOME/test_basic.log
|
||||
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m" >> "$HOME"/test_basic.log
|
||||
|
||||
# Run basic model test
|
||||
podman exec -it "$container_id" bash -c "
|
||||
@@ -43,7 +43,7 @@ function cpu_tests() {
|
||||
pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-google/gemma-1.1-2b-it]
|
||||
pytest -v -s tests/models/language/pooling/test_classification.py::test_models[float-jason9693/Qwen2.5-1.5B-apeach]
|
||||
# TODO: Below test case tests/models/language/pooling/test_embedding.py::test_models[True-ssmits/Qwen2-7B-Instruct-embed-base] fails on ppc64le. Disabling it for time being.
|
||||
# pytest -v -s tests/models/language/pooling/test_embedding.py -m cpu_model" >> $HOME/test_rest.log
|
||||
# pytest -v -s tests/models/language/pooling/test_embedding.py -m cpu_model" >> "$HOME"/test_rest.log
|
||||
}
|
||||
|
||||
# All of CPU tests are expected to be finished less than 40 mins.
|
||||
|
||||
@@ -16,5 +16,5 @@ echo "--- :docker: Building Docker image"
|
||||
docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
|
||||
|
||||
# Run the image, setting --shm-size=4g for tensor parallel.
|
||||
docker run --rm --cpuset-cpus=$CORE_RANGE --cpuset-mems=$NUMA_NODE -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g $IMAGE_NAME \
|
||||
timeout $TIMEOUT_VAL bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
|
||||
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
|
||||
timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
|
||||
|
||||
@@ -7,7 +7,7 @@ set -exuo pipefail
|
||||
# Try building the docker image
|
||||
image_name="hpu/upstream-vllm-ci:${BUILDKITE_COMMIT}"
|
||||
container_name="hpu-upstream-vllm-ci-${BUILDKITE_COMMIT}-container"
|
||||
cat <<EOF | docker build -t ${image_name} -f - .
|
||||
cat <<EOF | docker build -t "${image_name}" -f - .
|
||||
FROM gaudi-base-image:latest
|
||||
|
||||
COPY ./ /workspace/vllm
|
||||
@@ -39,12 +39,12 @@ EOF
|
||||
# functions, while other platforms only need one remove_docker_container
|
||||
# function.
|
||||
EXITCODE=1
|
||||
remove_docker_containers() { docker rm -f ${container_name} || true; }
|
||||
remove_docker_containers() { docker rm -f "${container_name}" || true; }
|
||||
trap 'remove_docker_containers; exit $EXITCODE;' EXIT
|
||||
remove_docker_containers
|
||||
|
||||
echo "Running HPU plugin v1 test"
|
||||
docker run --rm --runtime=habana --name=${container_name} --network=host \
|
||||
docker run --rm --runtime=habana --name="${container_name}" --network=host \
|
||||
-e HABANA_VISIBLE_DEVICES=all \
|
||||
-e VLLM_SKIP_WARMUP=true \
|
||||
-e PT_HPU_ENABLE_LAZY_COLLECTIVES=true \
|
||||
|
||||
@@ -41,6 +41,7 @@ get_config() {
|
||||
echo "Error: file '${TEST_RUN_CONFIG_FILE}' does not exist in the warehouse" >&2
|
||||
exit 1
|
||||
fi
|
||||
# shellcheck source=/dev/null
|
||||
source "${TEST_RUN_CONFIG_FILE}"
|
||||
echo "Base docker image name that get from configuration: ${BASE_IMAGE_NAME}"
|
||||
return 0
|
||||
@@ -48,9 +49,8 @@ get_config() {
|
||||
|
||||
# get test running configuration.
|
||||
fetch_vllm_test_cfg
|
||||
get_config
|
||||
# Check if the function call was successful. If not, exit the script.
|
||||
if [ $? -ne 0 ]; then
|
||||
if ! get_config; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
@@ -62,14 +62,14 @@ agent_idx=$(echo "${BUILDKITE_AGENT_NAME}" | awk -F'-' '{print $(NF-1)}')
|
||||
echo "agent_idx: ${agent_idx}"
|
||||
builder_name="cachebuilder${agent_idx}"
|
||||
builder_cache_dir="/mnt/docker-cache${agent_idx}"
|
||||
mkdir -p ${builder_cache_dir}
|
||||
mkdir -p "${builder_cache_dir}"
|
||||
|
||||
# Try building the docker image
|
||||
cat <<EOF | DOCKER_BUILDKIT=1 docker build \
|
||||
--add-host cache-service-vllm.nginx-pypi-cache.svc.cluster.local:${PYPI_CACHE_HOST} \
|
||||
--builder ${builder_name} --cache-from type=local,src=${builder_cache_dir} \
|
||||
--cache-to type=local,dest=${builder_cache_dir},mode=max \
|
||||
--progress=plain --load -t ${image_name} -f - .
|
||||
--add-host cache-service-vllm.nginx-pypi-cache.svc.cluster.local:"${PYPI_CACHE_HOST}" \
|
||||
--builder "${builder_name}" --cache-from type=local,src="${builder_cache_dir}" \
|
||||
--cache-to type=local,dest="${builder_cache_dir}",mode=max \
|
||||
--progress=plain --load -t "${image_name}" -f - .
|
||||
FROM ${BASE_IMAGE_NAME}
|
||||
|
||||
# Define environments
|
||||
@@ -116,7 +116,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/$(uname -i)-linux/devlib && \
|
||||
python3 -m pip install -v -e /workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
@@ -139,7 +139,7 @@ trap remove_docker_container EXIT
|
||||
# Generate corresponding --device args based on BUILDKITE_AGENT_NAME
|
||||
# Ascend NPU BUILDKITE_AGENT_NAME format is {hostname}-{agent_idx}-{npu_card_num}cards, and agent_idx starts from 1.
|
||||
# e.g. atlas-a2-001-1-2cards means this is the 1-th agent on atlas-a2-001 host, and it has 2 NPU cards.
|
||||
# returns --device /dev/davinci0 --device /dev/davinci1
|
||||
# returns one argument per line: --device, /dev/davinciX, ...
|
||||
parse_and_gen_devices() {
|
||||
local input="$1"
|
||||
local index cards_num
|
||||
@@ -151,29 +151,24 @@ parse_and_gen_devices() {
|
||||
return 1
|
||||
fi
|
||||
|
||||
local devices=""
|
||||
local i=0
|
||||
while (( i < cards_num )); do
|
||||
local dev_idx=$(((index - 1)*cards_num + i ))
|
||||
devices="$devices --device /dev/davinci${dev_idx}"
|
||||
printf '%s\n' "--device"
|
||||
printf '%s\n' "/dev/davinci${dev_idx}"
|
||||
((i++))
|
||||
done
|
||||
|
||||
# trim leading space
|
||||
devices="${devices#"${devices%%[![:space:]]*}"}"
|
||||
# Output devices: assigned to the caller variable
|
||||
printf '%s' "$devices"
|
||||
}
|
||||
|
||||
devices=$(parse_and_gen_devices "${BUILDKITE_AGENT_NAME}") || exit 1
|
||||
mapfile -t device_args < <(parse_and_gen_devices "${BUILDKITE_AGENT_NAME}") || exit 1
|
||||
|
||||
# Run the image and execute the Out-Of-Tree (OOT) platform interface test case on Ascend NPU hardware.
|
||||
# This test checks whether the OOT platform interface is functioning properly in conjunction with
|
||||
# the hardware plugin vllm-ascend.
|
||||
model_cache_dir=/mnt/modelscope${agent_idx}
|
||||
mkdir -p ${model_cache_dir}
|
||||
mkdir -p "${model_cache_dir}"
|
||||
docker run \
|
||||
${devices} \
|
||||
"${device_args[@]}" \
|
||||
--device /dev/davinci_manager \
|
||||
--device /dev/devmm_svm \
|
||||
--device /dev/hisi_hdc \
|
||||
@@ -182,7 +177,7 @@ docker run \
|
||||
-v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
|
||||
-v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
|
||||
-v /etc/ascend_install.info:/etc/ascend_install.info \
|
||||
-v ${model_cache_dir}:/root/.cache/modelscope \
|
||||
-v "${model_cache_dir}":/root/.cache/modelscope \
|
||||
--entrypoint="" \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
|
||||
@@ -61,7 +61,7 @@ echo "Results will be stored in: $RESULTS_DIR"
|
||||
echo "--- Installing Python dependencies ---"
|
||||
python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
|
||||
&& python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
|
||||
&& python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.9.2" \
|
||||
&& python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
|
||||
&& python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
|
||||
echo "--- Python dependencies installed ---"
|
||||
|
||||
|
||||
@@ -61,7 +61,7 @@ echo "Results will be stored in: $RESULTS_DIR"
|
||||
echo "--- Installing Python dependencies ---"
|
||||
python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
|
||||
&& python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
|
||||
&& python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.9.2" \
|
||||
&& python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
|
||||
&& python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
|
||||
echo "--- Python dependencies installed ---"
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ image_name="xpu/vllm-ci:${BUILDKITE_COMMIT}"
|
||||
container_name="xpu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
|
||||
|
||||
# Try building the docker image
|
||||
docker build -t ${image_name} -f docker/Dockerfile.xpu .
|
||||
docker build -t "${image_name}" -f docker/Dockerfile.xpu .
|
||||
|
||||
# Setup cleanup
|
||||
remove_docker_container() {
|
||||
|
||||
@@ -21,16 +21,16 @@ echo "Pushing original tag $ORIG_TAG_NAME$ORIG_TAG_SUFFIX to new nightly tag nam
|
||||
|
||||
# pull original arch-dependent images from AWS ECR Public
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-x86_64$ORIG_TAG_SUFFIX
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-aarch64$ORIG_TAG_SUFFIX
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-x86_64"$ORIG_TAG_SUFFIX"
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-aarch64"$ORIG_TAG_SUFFIX"
|
||||
# tag arch-dependent images
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-x86_64$ORIG_TAG_SUFFIX vllm/vllm-openai:$TAG_NAME-x86_64
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-aarch64$ORIG_TAG_SUFFIX vllm/vllm-openai:$TAG_NAME-aarch64
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-x86_64"$ORIG_TAG_SUFFIX" vllm/vllm-openai:"$TAG_NAME"-x86_64
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-aarch64"$ORIG_TAG_SUFFIX" vllm/vllm-openai:"$TAG_NAME"-aarch64
|
||||
# push arch-dependent images to DockerHub
|
||||
docker push vllm/vllm-openai:$TAG_NAME-x86_64
|
||||
docker push vllm/vllm-openai:$TAG_NAME-aarch64
|
||||
docker push vllm/vllm-openai:"$TAG_NAME"-x86_64
|
||||
docker push vllm/vllm-openai:"$TAG_NAME"-aarch64
|
||||
# push arch-independent manifest to DockerHub
|
||||
docker manifest create vllm/vllm-openai:$TAG_NAME vllm/vllm-openai:$TAG_NAME-x86_64 vllm/vllm-openai:$TAG_NAME-aarch64 --amend
|
||||
docker manifest create vllm/vllm-openai:$TAG_NAME-$BUILDKITE_COMMIT vllm/vllm-openai:$TAG_NAME-x86_64 vllm/vllm-openai:$TAG_NAME-aarch64 --amend
|
||||
docker manifest push vllm/vllm-openai:$TAG_NAME
|
||||
docker manifest push vllm/vllm-openai:$TAG_NAME-$BUILDKITE_COMMIT
|
||||
docker manifest create vllm/vllm-openai:"$TAG_NAME" vllm/vllm-openai:"$TAG_NAME"-x86_64 vllm/vllm-openai:"$TAG_NAME"-aarch64 --amend
|
||||
docker manifest create vllm/vllm-openai:"$TAG_NAME"-"$BUILDKITE_COMMIT" vllm/vllm-openai:"$TAG_NAME"-x86_64 vllm/vllm-openai:"$TAG_NAME"-aarch64 --amend
|
||||
docker manifest push vllm/vllm-openai:"$TAG_NAME"
|
||||
docker manifest push vllm/vllm-openai:"$TAG_NAME"-"$BUILDKITE_COMMIT"
|
||||
|
||||
@@ -67,7 +67,7 @@ start_nodes() {
|
||||
# 3. map the huggingface cache directory to the container
|
||||
# 3. assign ip addresses to the containers (head node: 192.168.10.10, worker nodes:
|
||||
# starting from 192.168.10.11)
|
||||
docker run -d $GPU_DEVICES --shm-size=10.24gb -e HF_TOKEN \
|
||||
docker run -d "$GPU_DEVICES" --shm-size=10.24gb -e HF_TOKEN \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface --name "node$node" \
|
||||
--network docker-net --ip 192.168.10.$((10 + $node)) --rm "$DOCKER_IMAGE" \
|
||||
/bin/bash -c "tail -f /dev/null"
|
||||
|
||||
@@ -29,7 +29,7 @@ fi
|
||||
if ! command -v uv &> /dev/null; then
|
||||
echo "Installing UV package manager..."
|
||||
curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
source $HOME/.local/bin/env
|
||||
source "$HOME"/.local/bin/env
|
||||
fi
|
||||
|
||||
# Clone Prime-RL repository at specific branch for reproducible tests
|
||||
|
||||
@@ -51,14 +51,14 @@ for BACK in "${BACKENDS[@]}"; do
|
||||
--enable-eplb \
|
||||
--trust-remote-code \
|
||||
--max-model-len 2048 \
|
||||
--all2all-backend $BACK \
|
||||
--port $PORT &
|
||||
--all2all-backend "$BACK" \
|
||||
--port "$PORT" &
|
||||
SERVER_PID=$!
|
||||
wait_for_server $PORT
|
||||
wait_for_server "$PORT"
|
||||
|
||||
TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
|
||||
OUT="${OUT_DIR}/${TAG}_${BACK}.json"
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port $PORT --num-questions ${NUM_Q} --save-results ${OUT}
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
|
||||
python3 - <<PY
|
||||
import json; acc=json.load(open('${OUT}'))['accuracy']
|
||||
print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")
|
||||
|
||||
@@ -47,20 +47,20 @@ for BACK in "${BACKENDS[@]}"; do
|
||||
vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--enable-eplb \
|
||||
--all2all-backend $BACK \
|
||||
--all2all-backend "$BACK" \
|
||||
--eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true}' \
|
||||
--tensor-parallel-size ${TENSOR_PARALLEL_SIZE} \
|
||||
--data-parallel-size ${DATA_PARALLEL_SIZE} \
|
||||
--tensor-parallel-size "${TENSOR_PARALLEL_SIZE}" \
|
||||
--data-parallel-size "${DATA_PARALLEL_SIZE}" \
|
||||
--enable-expert-parallel \
|
||||
--trust-remote-code \
|
||||
--max-model-len 2048 \
|
||||
--port $PORT &
|
||||
--port "$PORT" &
|
||||
SERVER_PID=$!
|
||||
wait_for_server $PORT
|
||||
wait_for_server "$PORT"
|
||||
|
||||
TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
|
||||
OUT="${OUT_DIR}/${TAG}_${BACK}.json"
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port $PORT --num-questions ${NUM_Q} --save-results ${OUT}
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
|
||||
python3 - <<PY
|
||||
import json; acc=json.load(open('${OUT}'))['accuracy']
|
||||
print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")
|
||||
|
||||
@@ -51,20 +51,20 @@ for BACK in "${BACKENDS[@]}"; do
|
||||
--tensor-parallel-size 4 \
|
||||
--enable-expert-parallel \
|
||||
--enable-eplb \
|
||||
--all2all-backend $BACK \
|
||||
--all2all-backend "$BACK" \
|
||||
--eplb-config '{"window_size":200,"step_interval":600,"use_async":true}' \
|
||||
--speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":1}' \
|
||||
--trust-remote-code \
|
||||
--max-model-len 2048 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
"${PLATFORM_ARGS[@]}" \
|
||||
--port $PORT &
|
||||
--port "$PORT" &
|
||||
SERVER_PID=$!
|
||||
wait_for_server $PORT
|
||||
wait_for_server "$PORT"
|
||||
|
||||
TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
|
||||
OUT="${OUT_DIR}/${TAG}_${BACK}.json"
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port $PORT --num-questions ${NUM_Q} --save-results ${OUT}
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
|
||||
python3 - <<PY
|
||||
import json; acc=json.load(open('${OUT}'))['accuracy']
|
||||
print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")
|
||||
|
||||
@@ -9,10 +9,11 @@ ENV_FILE=$1
|
||||
|
||||
# For testing on local vm, use `set -a` to export all variables
|
||||
source /etc/environment
|
||||
source $ENV_FILE
|
||||
# shellcheck source=/dev/null
|
||||
source "$ENV_FILE"
|
||||
|
||||
remove_docker_container() {
|
||||
docker rm -f $CONTAINER_NAME || true;
|
||||
docker rm -f "$CONTAINER_NAME" || true;
|
||||
}
|
||||
|
||||
trap remove_docker_container EXIT
|
||||
@@ -41,13 +42,13 @@ echo
|
||||
echo "starting docker...$CONTAINER_NAME"
|
||||
echo
|
||||
docker run \
|
||||
-v $DOWNLOAD_DIR:$DOWNLOAD_DIR \
|
||||
--env-file $ENV_FILE \
|
||||
-v "$DOWNLOAD_DIR":"$DOWNLOAD_DIR" \
|
||||
--env-file "$ENV_FILE" \
|
||||
-e HF_TOKEN="$HF_TOKEN" \
|
||||
-e TARGET_COMMIT=$BUILDKITE_COMMIT \
|
||||
-e MODEL=$MODEL \
|
||||
-e TARGET_COMMIT="$BUILDKITE_COMMIT" \
|
||||
-e MODEL="$MODEL" \
|
||||
-e WORKSPACE=/workspace \
|
||||
--name $CONTAINER_NAME \
|
||||
--name "$CONTAINER_NAME" \
|
||||
-d \
|
||||
--privileged \
|
||||
--network host \
|
||||
|
||||
@@ -42,21 +42,21 @@ echo "lanching vllm..."
|
||||
echo "logging to $VLLM_LOG"
|
||||
echo
|
||||
|
||||
vllm serve $MODEL \
|
||||
vllm serve "$MODEL" \
|
||||
--seed 42 \
|
||||
--max-num-seqs $MAX_NUM_SEQS \
|
||||
--max-num-batched-tokens $MAX_NUM_BATCHED_TOKENS \
|
||||
--tensor-parallel-size $TENSOR_PARALLEL_SIZE \
|
||||
--max-num-seqs "$MAX_NUM_SEQS" \
|
||||
--max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \
|
||||
--tensor-parallel-size "$TENSOR_PARALLEL_SIZE" \
|
||||
--no-enable-prefix-caching \
|
||||
--download_dir $DOWNLOAD_DIR \
|
||||
--max-model-len $MAX_MODEL_LEN > "$VLLM_LOG" 2>&1 &
|
||||
--download_dir "$DOWNLOAD_DIR" \
|
||||
--max-model-len "$MAX_MODEL_LEN" > "$VLLM_LOG" 2>&1 &
|
||||
|
||||
|
||||
echo "wait for 20 minutes.."
|
||||
echo
|
||||
# sleep 1200
|
||||
# wait for 10 minutes...
|
||||
for i in {1..120}; do
|
||||
for _ in {1..120}; do
|
||||
# TODO: detect other type of errors.
|
||||
if grep -Fq "raise RuntimeError" "$VLLM_LOG"; then
|
||||
echo "Detected RuntimeError, exiting."
|
||||
@@ -78,11 +78,11 @@ echo "logging to $BM_LOG"
|
||||
echo
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--dataset-name sonnet \
|
||||
--dataset-path benchmarks/sonnet_4x.txt \
|
||||
--sonnet-input-len $INPUT_LEN \
|
||||
--sonnet-output-len $OUTPUT_LEN \
|
||||
--sonnet-input-len "$INPUT_LEN" \
|
||||
--sonnet-output-len "$OUTPUT_LEN" \
|
||||
--ignore-eos > "$BM_LOG"
|
||||
|
||||
echo "completed..."
|
||||
|
||||
@@ -76,16 +76,15 @@ mkdir -p "$INDICES_OUTPUT_DIR"
|
||||
# this indices have relative paths that could work as long as it is next to the wheel directory in s3
|
||||
# i.e., the wheels are always in s3://vllm-wheels/<commit>/
|
||||
# and indices can be placed in /<commit>/, or /nightly/, or /<version>/
|
||||
if [[ ! -z "$DEFAULT_VARIANT_ALIAS" ]]; then
|
||||
alias_arg="--alias-to-default $DEFAULT_VARIANT_ALIAS"
|
||||
else
|
||||
alias_arg=""
|
||||
alias_args=()
|
||||
if [[ -n "$DEFAULT_VARIANT_ALIAS" ]]; then
|
||||
alias_args=(--alias-to-default "$DEFAULT_VARIANT_ALIAS")
|
||||
fi
|
||||
|
||||
# HACK: we do not need regex module here, but it is required by pre-commit hook
|
||||
# To avoid any external dependency, we simply replace it back to the stdlib re module
|
||||
sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
|
||||
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" $alias_arg
|
||||
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" "${alias_args[@]}"
|
||||
|
||||
# copy indices to /<commit>/ unconditionally
|
||||
echo "Uploading indices to $S3_COMMIT_PREFIX"
|
||||
@@ -100,9 +99,9 @@ fi
|
||||
# re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
|
||||
if [[ "$version" != *"dev"* ]]; then
|
||||
echo "Re-generating indices for /$pure_version/"
|
||||
rm -rf "$INDICES_OUTPUT_DIR/*"
|
||||
rm -rf "${INDICES_OUTPUT_DIR:?}/*"
|
||||
mkdir -p "$INDICES_OUTPUT_DIR"
|
||||
# wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
|
||||
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --wheel-dir "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" $alias_arg
|
||||
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --wheel-dir "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" "${alias_args[@]}"
|
||||
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/$pure_version/"
|
||||
fi
|
||||
|
||||
@@ -7,7 +7,7 @@ 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)
|
||||
GIT_VERSION=$(git describe --exact-match --tags "$BUILDKITE_COMMIT" 2>/dev/null)
|
||||
|
||||
echo "Release version from Buildkite: $RELEASE_VERSION"
|
||||
|
||||
@@ -55,7 +55,7 @@ mkdir -p $DIST_DIR
|
||||
aws s3 cp --recursive --exclude "*" --include "vllm-${PURE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" $DIST_DIR
|
||||
echo "Wheels copied to local directory"
|
||||
# generate source tarball
|
||||
git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" $BUILDKITE_COMMIT
|
||||
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)
|
||||
@@ -65,6 +65,6 @@ if [[ -z "$PYPI_WHEEL_FILES" ]]; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
python3 -m twine check $PYPI_WHEEL_FILES
|
||||
python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
|
||||
python3 -m twine check "$PYPI_WHEEL_FILES"
|
||||
python3 -m twine upload --non-interactive --verbose "$PYPI_WHEEL_FILES"
|
||||
echo "Wheels uploaded to PyPI"
|
||||
|
||||
@@ -55,7 +55,7 @@ mkdir -p all-rocm-wheels
|
||||
cp artifacts/rocm-base-wheels/*.whl all-rocm-wheels/ 2>/dev/null || true
|
||||
cp artifacts/rocm-vllm-wheel/*.whl all-rocm-wheels/ 2>/dev/null || true
|
||||
|
||||
WHEEL_COUNT=$(ls all-rocm-wheels/*.whl 2>/dev/null | wc -l)
|
||||
WHEEL_COUNT=$(find all-rocm-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
|
||||
echo "Total wheels to upload: $WHEEL_COUNT"
|
||||
|
||||
if [ "$WHEEL_COUNT" -eq 0 ]; then
|
||||
@@ -115,7 +115,7 @@ if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]] |
|
||||
fi
|
||||
|
||||
# Extract version from vLLM wheel and update version-specific index
|
||||
VLLM_WHEEL=$(ls all-rocm-wheels/vllm*.whl 2>/dev/null | head -1)
|
||||
VLLM_WHEEL=$(find all-rocm-wheels -maxdepth 1 -name 'vllm*.whl' 2>/dev/null | head -1)
|
||||
if [ -n "$VLLM_WHEEL" ]; then
|
||||
VERSION=$(unzip -p "$VLLM_WHEEL" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
|
||||
echo "Version in wheel: $VERSION"
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -197,6 +197,17 @@ steps:
|
||||
- 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: CrossLayer KV layout 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
|
||||
- CROSS_LAYERS_BLOCKS=True 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"
|
||||
|
||||
@@ -123,6 +123,7 @@ steps:
|
||||
- tests/test_inputs.py
|
||||
- tests/test_outputs.py
|
||||
- tests/test_pooling_params.py
|
||||
- tests/test_ray_env.py
|
||||
- tests/multimodal
|
||||
- tests/renderers
|
||||
- tests/standalone_tests/lazy_imports.py
|
||||
@@ -136,6 +137,7 @@ steps:
|
||||
- pytest -v -s test_inputs.py
|
||||
- pytest -v -s test_outputs.py
|
||||
- pytest -v -s test_pooling_params.py
|
||||
- pytest -v -s test_ray_env.py
|
||||
- pytest -v -s -m 'cpu_test' multimodal
|
||||
- pytest -v -s renderers
|
||||
- pytest -v -s tokenizers_
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
/build
|
||||
dist
|
||||
vllm/*.so
|
||||
.deps/
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
|
||||
+27
-11
@@ -2,7 +2,9 @@
|
||||
# for more info about CODEOWNERS file
|
||||
|
||||
# This lists cover the "core" components of vLLM that require careful review
|
||||
/vllm/executor/executor_base.py @zhuohan123 @youkaichao @alexm-redhat @njhill @22quinn
|
||||
/vllm/compilation @zou3519 @youkaichao @ProExpertProg
|
||||
/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery
|
||||
/vllm/lora @jeejeelee
|
||||
/vllm/model_executor/layers/attention @LucasWilkinson
|
||||
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
|
||||
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
|
||||
@@ -11,18 +13,34 @@
|
||||
/vllm/model_executor/layers/batch_invariant.py @yewentao256
|
||||
/vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
|
||||
/vllm/vllm_flash_attn @LucasWilkinson
|
||||
/vllm/lora @jeejeelee
|
||||
/vllm/reasoning @aarnphm @chaunceyjiang
|
||||
/vllm/entrypoints @aarnphm @chaunceyjiang
|
||||
/vllm/tool_parsers @aarnphm @chaunceyjiang
|
||||
/vllm/compilation @zou3519 @youkaichao @ProExpertProg
|
||||
/vllm/distributed/kv_transfer @NickLucche @ApostaC @orozery
|
||||
CMakeLists.txt @tlrmchlsmth @LucasWilkinson
|
||||
|
||||
# Any change to the VllmConfig changes can have a large user-facing impact,
|
||||
# so spam a lot of people
|
||||
/vllm/config @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg
|
||||
/vllm/config/cache.py @WoosukKwon @youkaichao @robertgshaw2-redhat @mgoin @tlrmchlsmth @houseroad @hmellor @yewentao256 @ProExpertProg @heheda12345
|
||||
/vllm/config/cache.py @heheda12345
|
||||
|
||||
# Entrypoints
|
||||
/vllm/entrypoints/anthropic @mgoin @DarkLight1337
|
||||
/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
|
||||
/vllm/entrypoints/mcp @heheda12345
|
||||
/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb
|
||||
/vllm/entrypoints/openai/realtime @njhill
|
||||
/vllm/entrypoints/openai/speech_to_text @NickLucche
|
||||
/vllm/entrypoints/pooling @noooop
|
||||
/vllm/entrypoints/sagemaker @DarkLight1337
|
||||
/vllm/entrypoints/serve @njhill
|
||||
/vllm/entrypoints/*.py @njhill
|
||||
/vllm/entrypoints/chat_utils.py @DarkLight1337
|
||||
/vllm/entrypoints/llm.py @DarkLight1337
|
||||
|
||||
# Input/Output Processing
|
||||
/vllm/sampling_params.py @njhill @NickLucche
|
||||
/vllm/pooling_params.py @noooop @DarkLight1337
|
||||
/vllm/tokenizers @DarkLight1337 @njhill
|
||||
/vllm/renderers @DarkLight1337 @njhill
|
||||
/vllm/reasoning @aarnphm @chaunceyjiang
|
||||
/vllm/tool_parsers @aarnphm @chaunceyjiang
|
||||
|
||||
# vLLM V1
|
||||
/vllm/v1/attention @LucasWilkinson
|
||||
@@ -115,8 +133,8 @@ mkdocs.yaml @hmellor
|
||||
/vllm/model_executor/models/mixtral*.py @patrickvonplaten
|
||||
/vllm/model_executor/models/voxtral*.py @patrickvonplaten
|
||||
/vllm/model_executor/models/pixtral*.py @patrickvonplaten
|
||||
/vllm/tokenizers/mistral.py @patrickvonplaten
|
||||
/vllm/transformers_utils/configs/mistral.py @patrickvonplaten
|
||||
/vllm/transformers_utils/tokenizers/mistral.py @patrickvonplaten
|
||||
|
||||
# Kernels
|
||||
/vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep
|
||||
@@ -152,9 +170,7 @@ mkdocs.yaml @hmellor
|
||||
/examples/pooling @noooop
|
||||
/tests/models/*/pooling* @noooop
|
||||
/tests/entrypoints/pooling @noooop
|
||||
/vllm/entrypoints/pooling @noooop
|
||||
/vllm/config/pooler.py @noooop
|
||||
/vllm/pooling_params.py @noooop
|
||||
/vllm/model_executor/layers/pooler @noooop
|
||||
|
||||
# Security guide and policies
|
||||
|
||||
@@ -38,7 +38,7 @@ repos:
|
||||
rev: 0.9.1
|
||||
hooks:
|
||||
- id: pip-compile
|
||||
args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu129, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
|
||||
args: [requirements/test.in, -o, requirements/test.txt, --index-strategy, unsafe-best-match, --torch-backend, cu128, --python-platform, x86_64-manylinux_2_28, --python-version, "3.12"]
|
||||
files: ^requirements/test\.(in|txt)$
|
||||
- repo: local
|
||||
hooks:
|
||||
|
||||
@@ -46,10 +46,10 @@ echo "VLLM_LOGGING_LEVEL=$VLLM_LOGGING_LEVEL"
|
||||
echo "RESULT_FILE=$RESULT"
|
||||
echo "====================== AUTO TUNEPARAMETERS ===================="
|
||||
|
||||
rm -rf $LOG_FOLDER
|
||||
rm -rf $PROFILE_PATH
|
||||
mkdir -p $LOG_FOLDER
|
||||
mkdir -p $PROFILE_PATH
|
||||
rm -rf "$LOG_FOLDER"
|
||||
rm -rf "$PROFILE_PATH"
|
||||
mkdir -p "$LOG_FOLDER"
|
||||
mkdir -p "$PROFILE_PATH"
|
||||
|
||||
cd "$BASE/vllm"
|
||||
|
||||
@@ -114,7 +114,7 @@ start_server() {
|
||||
|
||||
# wait for 10 minutes...
|
||||
server_started=0
|
||||
for i in {1..60}; do
|
||||
for _ in {1..60}; do
|
||||
# This line checks whether the server is still alive or not,
|
||||
# since that we should always have permission to send signal to the server process.
|
||||
kill -0 $server_pid 2> /dev/null || break
|
||||
@@ -145,12 +145,12 @@ run_benchmark() {
|
||||
local vllm_log="$LOG_FOLDER/vllm_log_${max_num_seqs}_${max_num_batched_tokens}.txt"
|
||||
echo "vllm_log: $vllm_log"
|
||||
echo
|
||||
rm -f $vllm_log
|
||||
rm -f "$vllm_log"
|
||||
pkill -if "vllm serve" || true
|
||||
|
||||
echo "starting server..."
|
||||
# Call start_server without a profile_dir to avoid profiling overhead
|
||||
start_server $gpu_memory_utilization $max_num_seqs $max_num_batched_tokens $vllm_log ""
|
||||
start_server "$gpu_memory_utilization" "$max_num_seqs" "$max_num_batched_tokens" "$vllm_log" ""
|
||||
result=$?
|
||||
if [[ "$result" -eq 1 ]]; then
|
||||
echo "server failed to start. gpu_memory_utilization:$gpu_memory_utilization, max_num_seqs:$max_num_seqs, max_num_batched_tokens: $max_num_batched_tokens"
|
||||
@@ -168,15 +168,15 @@ run_benchmark() {
|
||||
# --profile flag is removed from this call
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--dataset-name random \
|
||||
--random-input-len $adjusted_input_len \
|
||||
--random-output-len $OUTPUT_LEN \
|
||||
--random-output-len "$OUTPUT_LEN" \
|
||||
--ignore-eos \
|
||||
--disable-tqdm \
|
||||
--request-rate inf \
|
||||
--percentile-metrics ttft,tpot,itl,e2el \
|
||||
--goodput e2el:$MAX_LATENCY_ALLOWED_MS \
|
||||
--goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \
|
||||
--num-prompts 1000 \
|
||||
--random-prefix-len $prefix_len \
|
||||
--host "$HOSTNAME" \
|
||||
@@ -195,20 +195,20 @@ run_benchmark() {
|
||||
request_rate=$((${throughput%.*} + 1))
|
||||
while ((request_rate > 0)); do
|
||||
# clear prefix cache
|
||||
curl -X POST http://${HOSTNAME}:8004/reset_prefix_cache
|
||||
curl -X POST http://"${HOSTNAME}":8004/reset_prefix_cache
|
||||
sleep 5
|
||||
bm_log="$LOG_FOLDER/bm_log_${max_num_seqs}_${max_num_batched_tokens}_requestrate_${request_rate}.txt"
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--dataset-name random \
|
||||
--random-input-len $adjusted_input_len \
|
||||
--random-output-len $OUTPUT_LEN \
|
||||
--random-output-len "$OUTPUT_LEN" \
|
||||
--ignore-eos \
|
||||
--disable-tqdm \
|
||||
--request-rate $request_rate \
|
||||
--percentile-metrics ttft,tpot,itl,e2el \
|
||||
--goodput e2el:$MAX_LATENCY_ALLOWED_MS \
|
||||
--goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \
|
||||
--num-prompts 100 \
|
||||
--random-prefix-len $prefix_len \
|
||||
--host "$HOSTNAME" \
|
||||
@@ -255,7 +255,7 @@ gpu_memory_utilization=0.98
|
||||
find_gpu_memory_utilization=0
|
||||
while (( $(echo "$gpu_memory_utilization >= 0.9" | bc -l) )); do
|
||||
# Pass empty string for profile_dir argument
|
||||
start_server $gpu_memory_utilization "${num_seqs_list[-1]}" "${num_batched_tokens_list[-1]}" "$LOG_FOLDER/vllm_log_gpu_memory_utilization_$gpu_memory_utilization.log" ""
|
||||
start_server "$gpu_memory_utilization" "${num_seqs_list[-1]}" "${num_batched_tokens_list[-1]}" "$LOG_FOLDER/vllm_log_gpu_memory_utilization_$gpu_memory_utilization.log" ""
|
||||
result=$?
|
||||
if [[ "$result" -eq 0 ]]; then
|
||||
find_gpu_memory_utilization=1
|
||||
@@ -274,7 +274,7 @@ fi
|
||||
|
||||
for num_seqs in "${num_seqs_list[@]}"; do
|
||||
for num_batched_tokens in "${num_batched_tokens_list[@]}"; do
|
||||
run_benchmark $num_seqs $num_batched_tokens $gpu_memory_utilization
|
||||
run_benchmark "$num_seqs" "$num_batched_tokens" "$gpu_memory_utilization"
|
||||
done
|
||||
done
|
||||
echo "finish permutations"
|
||||
@@ -285,7 +285,7 @@ echo "finish permutations"
|
||||
if (( $(echo "$best_throughput > 0" | bc -l) )); then
|
||||
echo
|
||||
echo "Benchmark tuning finished. Now running profiling on the best configuration found..."
|
||||
echo "Best config: max_num_seqs: $best_max_num_seqs, max_num_batched_tokens: $best_num_batched_tokens, throughput: $best_throughput"
|
||||
echo "Best config: max_num_seqs: $best_max_num_seqs, max_num_batched_tokens: $best_num_batched_tokens, throughput: $best_throughput, goodput: $best_goodput"
|
||||
echo
|
||||
|
||||
vllm_log="$LOG_FOLDER/vllm_log_BEST_PROFILE.txt"
|
||||
@@ -293,7 +293,7 @@ if (( $(echo "$best_throughput > 0" | bc -l) )); then
|
||||
|
||||
# Start server with the best params and profiling ENABLED
|
||||
echo "Starting server for profiling..."
|
||||
start_server $gpu_memory_utilization $best_max_num_seqs $best_num_batched_tokens "$vllm_log" "$PROFILE_PATH"
|
||||
start_server "$gpu_memory_utilization" "$best_max_num_seqs" "$best_num_batched_tokens" "$vllm_log" "$PROFILE_PATH"
|
||||
|
||||
# Run benchmark with the best params and the --profile flag
|
||||
echo "Running benchmark with profiling..."
|
||||
@@ -301,15 +301,15 @@ if (( $(echo "$best_throughput > 0" | bc -l) )); then
|
||||
adjusted_input_len=$(( INPUT_LEN - prefix_len ))
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--dataset-name random \
|
||||
--random-input-len $adjusted_input_len \
|
||||
--random-output-len $OUTPUT_LEN \
|
||||
--random-output-len "$OUTPUT_LEN" \
|
||||
--ignore-eos \
|
||||
--disable-tqdm \
|
||||
--request-rate $best_request_rate \
|
||||
--request-rate "$best_request_rate" \
|
||||
--percentile-metrics ttft,tpot,itl,e2el \
|
||||
--goodput e2el:$MAX_LATENCY_ALLOWED_MS \
|
||||
--goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \
|
||||
--num-prompts 100 \
|
||||
--random-prefix-len $prefix_len \
|
||||
--host "$HOSTNAME" \
|
||||
|
||||
@@ -64,7 +64,7 @@ for i in $(seq 0 $(($num_runs - 1))); do
|
||||
else
|
||||
STATUS="FAILURE"
|
||||
((FAILURE_COUNT++))
|
||||
FAILED_RUNS+=("Run #$((i+1)): $(echo $run_object | jq -c .)")
|
||||
FAILED_RUNS+=("Run #$((i+1)): $(echo "$run_object" | jq -c .)")
|
||||
fi
|
||||
|
||||
RUN_OUTPUT=$(<"$RUN_OUTPUT_FILE")
|
||||
|
||||
@@ -0,0 +1,471 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations.
|
||||
|
||||
Compares:
|
||||
- apply_top_k_top_p_triton (Triton binary search)
|
||||
- apply_top_k_top_p (PyTorch sort-based)
|
||||
|
||||
Scenarios:
|
||||
- top_k only (whole batch, partial batch)
|
||||
- top_p only (whole batch, partial batch)
|
||||
- mix of top_k and top_p
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch
|
||||
from vllm.v1.sample.ops.topk_topp_triton import (
|
||||
apply_top_k_top_p_triton,
|
||||
reset_buffer_cache,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkConfig:
|
||||
"""Configuration for a benchmark run."""
|
||||
|
||||
name: str
|
||||
batch_size: int
|
||||
vocab_size: int
|
||||
# k and p can be tensors or None
|
||||
k_values: torch.Tensor | None # [batch_size] or None
|
||||
p_values: torch.Tensor | None # [batch_size] or None
|
||||
description: str
|
||||
ops_pct: float = 0.0 # Percentage of ops relative to batch size
|
||||
|
||||
|
||||
def calculate_ops_pct(
|
||||
k_values: torch.Tensor | None,
|
||||
p_values: torch.Tensor | None,
|
||||
vocab_size: int,
|
||||
batch_size: int,
|
||||
) -> float:
|
||||
"""
|
||||
Calculate the percentage of active top-k and top-p operations.
|
||||
|
||||
Returns percentage where 100% = batch_size ops.
|
||||
E.g., if all rows have both top-k and top-p active, returns 200%.
|
||||
"""
|
||||
active_ops = 0
|
||||
|
||||
if k_values is not None:
|
||||
# Count rows where k < vocab_size (active top-k filtering)
|
||||
active_ops += (k_values < vocab_size).sum().item()
|
||||
|
||||
if p_values is not None:
|
||||
# Count rows where p < 1.0 (active top-p filtering)
|
||||
active_ops += (p_values < 1.0).sum().item()
|
||||
|
||||
return (active_ops / batch_size) * 100 if batch_size > 0 else 0.0
|
||||
|
||||
|
||||
def create_logits(
|
||||
batch_size: int, vocab_size: int, device: str = "cuda"
|
||||
) -> torch.Tensor:
|
||||
"""Create random logits mimicking a realistic LLM distribution.
|
||||
|
||||
Uses a Zipf-like probability distribution (rank^-1.1) converted to logits
|
||||
via log, then randomly permuted per row. This produces a peaked distribution
|
||||
where a small number of tokens capture most probability mass, similar to
|
||||
real model outputs.
|
||||
"""
|
||||
# Create Zipf-like probabilities: p(rank) ~ rank^(-alpha)
|
||||
ranks = torch.arange(1, vocab_size + 1, dtype=torch.float32, device=device)
|
||||
probs = ranks.pow(-1.1)
|
||||
probs = probs / probs.sum()
|
||||
|
||||
# Convert to logits (log-probabilities, unnormalized is fine)
|
||||
base_logits = probs.log()
|
||||
|
||||
# Broadcast to batch and randomly permute each row
|
||||
logits = base_logits.unsqueeze(0).expand(batch_size, -1).clone()
|
||||
for i in range(batch_size):
|
||||
logits[i] = logits[i, torch.randperm(vocab_size, device=device)]
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
def measure_memory() -> tuple[int, int]:
|
||||
"""Return (allocated, reserved) memory in bytes."""
|
||||
torch.cuda.synchronize()
|
||||
return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()
|
||||
|
||||
|
||||
def reset_memory_stats():
|
||||
"""Reset peak memory statistics."""
|
||||
reset_buffer_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
def benchmark_function(
|
||||
func,
|
||||
logits: torch.Tensor,
|
||||
k: torch.Tensor | None,
|
||||
p: torch.Tensor | None,
|
||||
warmup_iters: int = 5,
|
||||
benchmark_iters: int = 20,
|
||||
) -> tuple[float, int]:
|
||||
"""
|
||||
Benchmark a function and return (avg_time_ms, peak_memory_bytes).
|
||||
|
||||
Returns average time in milliseconds and peak memory usage.
|
||||
"""
|
||||
# Warmup
|
||||
for _ in range(warmup_iters):
|
||||
logits_copy = logits.clone()
|
||||
func(logits_copy, k, p)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Reset memory stats before benchmark
|
||||
reset_memory_stats()
|
||||
|
||||
# Benchmark
|
||||
start_events = [
|
||||
torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)
|
||||
]
|
||||
end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)]
|
||||
|
||||
for i in range(benchmark_iters):
|
||||
logits_copy = logits.clone()
|
||||
start_events[i].record()
|
||||
func(logits_copy, k, p)
|
||||
end_events[i].record()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Calculate timing
|
||||
times = [
|
||||
start_events[i].elapsed_time(end_events[i]) for i in range(benchmark_iters)
|
||||
]
|
||||
avg_time = sum(times) / len(times)
|
||||
|
||||
# Get peak memory
|
||||
_, peak_memory = measure_memory()
|
||||
|
||||
return avg_time, peak_memory
|
||||
|
||||
|
||||
def create_benchmark_configs(
|
||||
batch_sizes: list[int],
|
||||
vocab_sizes: list[int],
|
||||
device: str = "cuda",
|
||||
) -> list[BenchmarkConfig]:
|
||||
"""Create all benchmark configurations."""
|
||||
configs = []
|
||||
|
||||
for vocab_size in vocab_sizes:
|
||||
for batch_size in batch_sizes:
|
||||
# 1. Top-k only - whole batch (all rows have k < vocab_size)
|
||||
k_all = torch.full((batch_size,), 50, dtype=torch.int32, device=device)
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"topk_whole_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=k_all,
|
||||
p_values=None,
|
||||
description=f"Top-k only (whole batch, k=50), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(k_all, None, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
# 2. Top-k only - partial batch (half have k=50, half have k=vocab_size)
|
||||
k_partial = torch.full((batch_size,), 50, dtype=torch.int32, device=device)
|
||||
k_partial[batch_size // 2 :] = vocab_size # No filtering for second half
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"topk_partial_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=k_partial,
|
||||
p_values=None,
|
||||
description=f"Top-k only (partial batch, 50% k=50, 50% k=vocab), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(k_partial, None, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
# 3. Top-p only - whole batch (all rows have p < 1.0)
|
||||
p_all = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device)
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"topp_whole_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=None,
|
||||
p_values=p_all,
|
||||
description=f"Top-p only (whole batch, p=0.9), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(None, p_all, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
# 4. Top-p only - partial batch (half have p=0.9, half have p=1.0)
|
||||
p_partial = torch.full(
|
||||
(batch_size,), 0.9, dtype=torch.float32, device=device
|
||||
)
|
||||
p_partial[batch_size // 2 :] = 1.0 # No filtering for second half
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"topp_partial_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=None,
|
||||
p_values=p_partial,
|
||||
description=f"Top-p only (partial batch, 50% p=0.9, 50% p=1.0), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(None, p_partial, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
# 5. Mix of top-k and top-p (both applied to whole batch)
|
||||
k_mix = torch.full((batch_size,), 100, dtype=torch.int32, device=device)
|
||||
p_mix = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device)
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"topk_topp_whole_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=k_mix,
|
||||
p_values=p_mix,
|
||||
description=f"Top-k + Top-p (whole batch, k=100, p=0.9), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(k_mix, p_mix, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
# 6. Mix with partial application (some rows k only, some p only, some both)
|
||||
k_mixed = torch.full(
|
||||
(batch_size,), vocab_size, dtype=torch.int32, device=device
|
||||
)
|
||||
p_mixed = torch.full((batch_size,), 1.0, dtype=torch.float32, device=device)
|
||||
# First third: k only
|
||||
third = batch_size // 3
|
||||
k_mixed[:third] = 50
|
||||
# Second third: p only
|
||||
p_mixed[third : 2 * third] = 0.5
|
||||
# Last third: both k and p
|
||||
k_mixed[2 * third :] = 100
|
||||
p_mixed[2 * third :] = 0.9
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"mixed_partial_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=k_mixed,
|
||||
p_values=p_mixed,
|
||||
description=f"Mixed partial (1/3 k=50, 1/3 p=0.9, 1/3 both), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(k_mixed, p_mixed, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
return configs
|
||||
|
||||
|
||||
def format_memory(bytes_val: int) -> str:
|
||||
"""Format memory in human-readable form."""
|
||||
if bytes_val >= 1024**3:
|
||||
return f"{bytes_val / (1024**3):.2f} GB"
|
||||
elif bytes_val >= 1024**2:
|
||||
return f"{bytes_val / (1024**2):.2f} MB"
|
||||
elif bytes_val >= 1024:
|
||||
return f"{bytes_val / 1024:.2f} KB"
|
||||
return f"{bytes_val} B"
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
configs: list[BenchmarkConfig],
|
||||
warmup_iters: int = 5,
|
||||
benchmark_iters: int = 20,
|
||||
verbose: bool = True,
|
||||
):
|
||||
"""Run all benchmarks and print results."""
|
||||
results = []
|
||||
|
||||
print("=" * 100)
|
||||
print("Top-k/Top-p Benchmark: Triton vs PyTorch Sort-based")
|
||||
print("=" * 100)
|
||||
print()
|
||||
|
||||
for config in configs:
|
||||
if verbose:
|
||||
print(f"Running: {config.description}")
|
||||
|
||||
# Create fresh logits for this config
|
||||
logits = create_logits(config.batch_size, config.vocab_size)
|
||||
|
||||
# Benchmark Triton
|
||||
reset_memory_stats()
|
||||
triton_time, triton_mem = benchmark_function(
|
||||
apply_top_k_top_p_triton,
|
||||
logits,
|
||||
config.k_values,
|
||||
config.p_values,
|
||||
warmup_iters,
|
||||
benchmark_iters,
|
||||
)
|
||||
|
||||
# Benchmark PyTorch
|
||||
reset_memory_stats()
|
||||
pytorch_time, pytorch_mem = benchmark_function(
|
||||
apply_top_k_top_p_pytorch,
|
||||
logits,
|
||||
config.k_values,
|
||||
config.p_values,
|
||||
warmup_iters,
|
||||
benchmark_iters,
|
||||
)
|
||||
|
||||
speedup = pytorch_time / triton_time if triton_time > 0 else float("inf")
|
||||
mem_ratio = pytorch_mem / triton_mem if triton_mem > 0 else float("inf")
|
||||
|
||||
result = {
|
||||
"config": config,
|
||||
"triton_time_ms": triton_time,
|
||||
"pytorch_time_ms": pytorch_time,
|
||||
"triton_mem": triton_mem,
|
||||
"pytorch_mem": pytorch_mem,
|
||||
"speedup": speedup,
|
||||
"mem_ratio": mem_ratio,
|
||||
}
|
||||
results.append(result)
|
||||
|
||||
if verbose:
|
||||
print(f" Triton: {triton_time:.3f} ms, {format_memory(triton_mem)}")
|
||||
print(f" PyTorch: {pytorch_time:.3f} ms, {format_memory(pytorch_mem)}")
|
||||
print(f" Speedup: {speedup:.2f}x, Memory ratio: {mem_ratio:.2f}x")
|
||||
print()
|
||||
|
||||
# Clean up
|
||||
del logits
|
||||
reset_memory_stats()
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def print_summary_table(results: list[dict]):
|
||||
"""Print a summary table of results."""
|
||||
print()
|
||||
print("=" * 130)
|
||||
print("SUMMARY TABLE")
|
||||
print("=" * 130)
|
||||
print()
|
||||
|
||||
# Header
|
||||
header = (
|
||||
f"{'Scenario':<40} {'Batch':>6} {'Vocab':>7} {'Ops%':>6} "
|
||||
f"{'Triton (ms)':>12} {'PyTorch (ms)':>13} {'Speedup':>8} "
|
||||
f"{'Tri Mem':>10} {'Pyt Mem':>10}"
|
||||
)
|
||||
print(header)
|
||||
print("-" * 130)
|
||||
|
||||
# Group by scenario type
|
||||
current_vocab = None
|
||||
for result in results:
|
||||
config = result["config"]
|
||||
|
||||
# Add separator between vocab sizes
|
||||
if current_vocab != config.vocab_size:
|
||||
if current_vocab is not None:
|
||||
print("-" * 130)
|
||||
current_vocab = config.vocab_size
|
||||
|
||||
scenario = config.name.split("_b")[0] # Extract scenario name
|
||||
print(
|
||||
f"{scenario:<40} {config.batch_size:>6} {config.vocab_size:>7} "
|
||||
f"{config.ops_pct:>5.0f}% "
|
||||
f"{result['triton_time_ms']:>12.3f} {result['pytorch_time_ms']:>13.3f} "
|
||||
f"{result['speedup']:>7.2f}x "
|
||||
f"{format_memory(result['triton_mem']):>10} "
|
||||
f"{format_memory(result['pytorch_mem']):>10}"
|
||||
)
|
||||
|
||||
print("=" * 130)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark Triton vs PyTorch sort-based top-k/top-p implementations"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-sizes",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[1, 4, 16, 64, 128, 512, 1024, 2048],
|
||||
help="Batch sizes to test (default: 1 4 16 64)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vocab-sizes",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[32768, 131072], # 32k, 128k
|
||||
help="Vocabulary sizes to test (default: 32768 131072)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-iters",
|
||||
type=int,
|
||||
default=5,
|
||||
help="Number of warmup iterations (default: 5)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--benchmark-iters",
|
||||
type=int,
|
||||
default=20,
|
||||
help="Number of benchmark iterations (default: 20)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quiet",
|
||||
action="store_true",
|
||||
help="Only print summary table",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Print configuration
|
||||
print(f"Batch sizes: {args.batch_sizes}")
|
||||
print(f"Vocab sizes: {args.vocab_sizes}")
|
||||
print(f"Warmup iterations: {args.warmup_iters}")
|
||||
print(f"Benchmark iterations: {args.benchmark_iters}")
|
||||
print()
|
||||
|
||||
# Check CUDA
|
||||
if not torch.cuda.is_available():
|
||||
print("ERROR: CUDA is not available. This benchmark requires a GPU.")
|
||||
return
|
||||
|
||||
device_name = torch.cuda.get_device_name(0)
|
||||
print(f"GPU: {device_name}")
|
||||
print()
|
||||
|
||||
# Create configs
|
||||
configs = create_benchmark_configs(
|
||||
args.batch_sizes,
|
||||
args.vocab_sizes,
|
||||
)
|
||||
|
||||
# Run benchmarks
|
||||
results = run_benchmark(
|
||||
configs,
|
||||
warmup_iters=args.warmup_iters,
|
||||
benchmark_iters=args.benchmark_iters,
|
||||
verbose=not args.quiet,
|
||||
)
|
||||
|
||||
# Print summary
|
||||
print_summary_table(results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -100,13 +100,38 @@ def benchmark_config(
|
||||
dtype: torch.dtype,
|
||||
use_fp8_w8a8: bool,
|
||||
use_int8_w8a16: bool,
|
||||
use_int4_w4a16: bool = False,
|
||||
num_iters: int = 100,
|
||||
block_quant_shape: list[int] = None,
|
||||
use_deep_gemm: bool = False,
|
||||
) -> float:
|
||||
init_dtype = torch.float16 if use_fp8_w8a8 else dtype
|
||||
x = torch.randn(num_tokens, hidden_size, dtype=dtype)
|
||||
if use_int8_w8a16:
|
||||
if use_int4_w4a16:
|
||||
# Int4 packed weights: 2 int4 values per uint8 byte
|
||||
# K dimension is packed (halved)
|
||||
intermediate_size = shard_intermediate_size // 2 # after silu_and_mul
|
||||
w1 = torch.randint(
|
||||
0,
|
||||
255,
|
||||
(
|
||||
num_experts,
|
||||
shard_intermediate_size,
|
||||
hidden_size // 2, # int4 packing
|
||||
),
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
w2 = torch.randint(
|
||||
0,
|
||||
255,
|
||||
(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
intermediate_size // 2, # int4 packing
|
||||
),
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
elif use_int8_w8a16:
|
||||
w1 = torch.randint(
|
||||
-127,
|
||||
127,
|
||||
@@ -140,7 +165,20 @@ def benchmark_config(
|
||||
w2_scale = None
|
||||
a1_scale = None
|
||||
a2_scale = None
|
||||
if use_int8_w8a16:
|
||||
if use_int4_w4a16:
|
||||
if block_quant_shape is None:
|
||||
raise ValueError("block_quant_shape is required for int4_w4a16")
|
||||
group_size = block_quant_shape[1]
|
||||
# Scales shape: (E, N, K // group_size) in fp16
|
||||
w1_scale = torch.rand(
|
||||
(num_experts, shard_intermediate_size, hidden_size // group_size),
|
||||
dtype=dtype,
|
||||
)
|
||||
w2_scale = torch.rand(
|
||||
(num_experts, hidden_size, intermediate_size // group_size),
|
||||
dtype=dtype,
|
||||
)
|
||||
elif use_int8_w8a16:
|
||||
w1_scale = torch.randn(
|
||||
(num_experts, 2 * shard_intermediate_size), dtype=torch.float32
|
||||
)
|
||||
@@ -199,6 +237,7 @@ def benchmark_config(
|
||||
a1_scale=a1_scale,
|
||||
a2_scale=a2_scale,
|
||||
block_shape=block_quant_shape,
|
||||
weight_dtype="int4" if use_int4_w4a16 else None,
|
||||
)
|
||||
|
||||
deep_gemm_experts = None
|
||||
@@ -481,6 +520,7 @@ class BenchmarkWorker:
|
||||
dtype: torch.dtype,
|
||||
use_fp8_w8a8: bool,
|
||||
use_int8_w8a16: bool,
|
||||
use_int4_w4a16: bool = False,
|
||||
block_quant_shape: list[int] = None,
|
||||
use_deep_gemm: bool = False,
|
||||
) -> tuple[dict[str, int], float]:
|
||||
@@ -488,7 +528,10 @@ class BenchmarkWorker:
|
||||
|
||||
set_random_seed(self.seed)
|
||||
dtype_str = _get_config_dtype_str(
|
||||
dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
|
||||
dtype,
|
||||
use_int8_w8a16=use_int8_w8a16,
|
||||
use_fp8_w8a8=use_fp8_w8a8,
|
||||
use_int4_w4a16=use_int4_w4a16,
|
||||
)
|
||||
# NOTE(woosuk): The current naming convention uses w2.shape[2], which
|
||||
# is the intermediate size after silu_and_mul.
|
||||
@@ -519,6 +562,7 @@ class BenchmarkWorker:
|
||||
dtype,
|
||||
use_fp8_w8a8,
|
||||
use_int8_w8a16,
|
||||
use_int4_w4a16=use_int4_w4a16,
|
||||
num_iters=100,
|
||||
block_quant_shape=block_quant_shape,
|
||||
use_deep_gemm=use_deep_gemm,
|
||||
@@ -535,6 +579,7 @@ class BenchmarkWorker:
|
||||
dtype: torch.dtype,
|
||||
use_fp8_w8a8: bool,
|
||||
use_int8_w8a16: bool,
|
||||
use_int4_w4a16: bool,
|
||||
search_space: list[dict[str, int]],
|
||||
block_quant_shape: list[int],
|
||||
use_deep_gemm: bool,
|
||||
@@ -545,7 +590,7 @@ class BenchmarkWorker:
|
||||
best_config = None
|
||||
best_time = float("inf")
|
||||
if current_platform.is_rocm():
|
||||
is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)
|
||||
is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16 or use_int4_w4a16)
|
||||
search_space = prune_rocm_search_space(
|
||||
num_tokens,
|
||||
shard_intermediate_size,
|
||||
@@ -574,6 +619,7 @@ class BenchmarkWorker:
|
||||
dtype,
|
||||
use_fp8_w8a8,
|
||||
use_int8_w8a16,
|
||||
use_int4_w4a16,
|
||||
num_iters=20,
|
||||
block_quant_shape=block_quant_shape,
|
||||
use_deep_gemm=use_deep_gemm,
|
||||
@@ -621,6 +667,7 @@ def sort_config(config: BenchmarkConfig) -> BenchmarkConfig:
|
||||
else {}
|
||||
),
|
||||
**({"kpack": config["kpack"]} if "kpack" in config else {}),
|
||||
**({"SPLIT_K": config["SPLIT_K"]} if "SPLIT_K" in config else {}),
|
||||
}
|
||||
|
||||
|
||||
@@ -633,11 +680,15 @@ def save_configs(
|
||||
dtype: torch.dtype,
|
||||
use_fp8_w8a8: bool,
|
||||
use_int8_w8a16: bool,
|
||||
use_int4_w4a16: bool,
|
||||
block_quant_shape: list[int],
|
||||
save_dir: str,
|
||||
) -> None:
|
||||
dtype_str = _get_config_dtype_str(
|
||||
dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
|
||||
dtype,
|
||||
use_int8_w8a16=use_int8_w8a16,
|
||||
use_fp8_w8a8=use_fp8_w8a8,
|
||||
use_int4_w4a16=use_int4_w4a16,
|
||||
)
|
||||
|
||||
# NOTE(woosuk): The current naming convention uses w2.shape[2], which
|
||||
@@ -739,6 +790,38 @@ def get_model_params(config):
|
||||
return E, topk, intermediate_size, hidden_size
|
||||
|
||||
|
||||
def get_quantization_group_size(config) -> int | None:
|
||||
"""Extract the quantization group size from the HF model config.
|
||||
|
||||
This reads directly from the HuggingFace config object (as returned by
|
||||
``get_config()``), not from vLLM's quantization config classes.
|
||||
|
||||
Supports AWQ/GPTQ-style configs (direct 'group_size' key) and
|
||||
compressed-tensors configs (nested inside 'config_groups').
|
||||
"""
|
||||
quantization_config = getattr(config, "quantization_config", {})
|
||||
if not isinstance(quantization_config, dict):
|
||||
return None
|
||||
# AWQ / GPTQ style: group_size is a top-level key
|
||||
gs = quantization_config.get("group_size")
|
||||
if gs is not None:
|
||||
return gs
|
||||
# compressed-tensors style: group_size is nested in config_groups
|
||||
config_groups = quantization_config.get("config_groups", {})
|
||||
if not isinstance(config_groups, dict):
|
||||
return None
|
||||
for group_cfg in config_groups.values():
|
||||
if not isinstance(group_cfg, dict):
|
||||
continue
|
||||
weights = group_cfg.get("weights", {})
|
||||
if not isinstance(weights, dict):
|
||||
continue
|
||||
gs = weights.get("group_size")
|
||||
if gs is not None:
|
||||
return gs
|
||||
return None
|
||||
|
||||
|
||||
def main(args: argparse.Namespace):
|
||||
print(args)
|
||||
|
||||
@@ -757,7 +840,20 @@ 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_int4_w4a16 = args.dtype == "int4_w4a16"
|
||||
block_quant_shape = get_weight_block_size_safety(config)
|
||||
if use_int4_w4a16:
|
||||
group_size = get_quantization_group_size(config)
|
||||
if group_size is None:
|
||||
raise ValueError(
|
||||
"Could not determine group_size from model config. "
|
||||
"The model's quantization_config must contain a 'group_size' "
|
||||
"field (AWQ/GPTQ) or 'config_groups.*.weights.group_size' "
|
||||
"(compressed-tensors)."
|
||||
)
|
||||
# For int4_w4a16, block_shape = [0, group_size]
|
||||
# block_shape[0]=0 means no block quantization on N dimension
|
||||
block_quant_shape = [0, group_size]
|
||||
|
||||
if args.batch_size is None:
|
||||
batch_sizes = [
|
||||
@@ -811,8 +907,20 @@ def main(args: argparse.Namespace):
|
||||
return ray.get(outputs)
|
||||
|
||||
if args.tune:
|
||||
is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)
|
||||
search_space = get_configs_compute_bound(is_fp16, block_quant_shape)
|
||||
# int4_w4a16 weights are uint8-packed, not fp16; treat like fp8 for
|
||||
# search space generation (no matrix_instr_nonkdim/kpack exploration).
|
||||
is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16 or use_int4_w4a16)
|
||||
# For int4_w4a16, the group_size constraint on BLOCK_SIZE_K does not
|
||||
# apply: the gptq_awq kernel handles arbitrary BLOCK_SIZE_K regardless
|
||||
# of group_size. Skip block_quant_shape filtering to keep the full
|
||||
# search space (e.g. BLOCK_SIZE_K=64 with group_size=128).
|
||||
tune_block_quant_shape = None if use_int4_w4a16 else block_quant_shape
|
||||
search_space = get_configs_compute_bound(is_fp16, tune_block_quant_shape)
|
||||
if use_int4_w4a16:
|
||||
# SPLIT_K is a required kernel constexpr for gptq_awq kernel;
|
||||
# only SPLIT_K=1 is used at runtime, so fix it during tuning.
|
||||
for cfg in search_space:
|
||||
cfg["SPLIT_K"] = 1
|
||||
print(f"Start tuning over {len(search_space)} configurations...")
|
||||
if use_deep_gemm:
|
||||
raise ValueError(
|
||||
@@ -832,6 +940,7 @@ def main(args: argparse.Namespace):
|
||||
dtype,
|
||||
use_fp8_w8a8,
|
||||
use_int8_w8a16,
|
||||
use_int4_w4a16,
|
||||
search_space,
|
||||
block_quant_shape,
|
||||
use_deep_gemm,
|
||||
@@ -851,6 +960,7 @@ def main(args: argparse.Namespace):
|
||||
dtype,
|
||||
use_fp8_w8a8,
|
||||
use_int8_w8a16,
|
||||
use_int4_w4a16,
|
||||
block_quant_shape,
|
||||
args.save_dir,
|
||||
)
|
||||
@@ -869,6 +979,7 @@ def main(args: argparse.Namespace):
|
||||
dtype,
|
||||
use_fp8_w8a8,
|
||||
use_int8_w8a16,
|
||||
use_int4_w4a16,
|
||||
block_quant_shape,
|
||||
use_deep_gemm,
|
||||
)
|
||||
@@ -891,7 +1002,10 @@ if __name__ == "__main__":
|
||||
)
|
||||
parser.add_argument("--enable-expert-parallel", "-enable-ep", action="store_true")
|
||||
parser.add_argument(
|
||||
"--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"
|
||||
"--dtype",
|
||||
type=str,
|
||||
choices=["auto", "fp8_w8a8", "int8_w8a16", "int4_w4a16"],
|
||||
default="auto",
|
||||
)
|
||||
parser.add_argument("--use-deep-gemm", action="store_true")
|
||||
parser.add_argument(
|
||||
|
||||
@@ -71,7 +71,7 @@ while [[ $# -gt 0 ]]; do
|
||||
usage
|
||||
;;
|
||||
*)
|
||||
echo "Unknown argument: $1\n"
|
||||
printf "Unknown argument: %s\n" "$1"
|
||||
usage
|
||||
;;
|
||||
esac
|
||||
@@ -84,15 +84,17 @@ mkdir -p "$OUTPUT_DIR"
|
||||
QPS_VALUES=(25 20 15 10 5 1)
|
||||
|
||||
# Common parameters
|
||||
COMMON_PARAMS="--backend $BACKEND \
|
||||
--model $MODEL \
|
||||
--dataset $DATASET \
|
||||
--structured-output-ratio $STRUCTURED_OUTPUT_RATIO \
|
||||
--save-results \
|
||||
--result-dir $OUTPUT_DIR \
|
||||
--output-len $MAX_NEW_TOKENS \
|
||||
--port $PORT \
|
||||
--tokenizer-mode $TOKENIZER_MODE"
|
||||
COMMON_PARAMS=(
|
||||
--backend "$BACKEND"
|
||||
--model "$MODEL"
|
||||
--dataset "$DATASET"
|
||||
--structured-output-ratio "$STRUCTURED_OUTPUT_RATIO"
|
||||
--save-results
|
||||
--result-dir "$OUTPUT_DIR"
|
||||
--output-len "$MAX_NEW_TOKENS"
|
||||
--port "$PORT"
|
||||
--tokenizer-mode "$TOKENIZER_MODE"
|
||||
)
|
||||
|
||||
echo "Starting structured output benchmark with model: $MODEL"
|
||||
echo "Backend: $BACKEND"
|
||||
@@ -109,17 +111,17 @@ for qps in "${QPS_VALUES[@]}"; do
|
||||
GIT_BRANCH=$(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo "unknown")
|
||||
|
||||
# Construct filename for this run
|
||||
FILENAME="${BACKEND}_${qps}qps_$(basename $MODEL)_${DATASET}_${GIT_HASH}.json"
|
||||
FILENAME="${BACKEND}_${qps}qps_$(basename "$MODEL")_${DATASET}_${GIT_HASH}_${GIT_BRANCH}.json"
|
||||
|
||||
NUM_PROMPTS=$(echo "$TOTAL_SECONDS * $qps" | bc)
|
||||
NUM_PROMPTS=${NUM_PROMPTS%.*} # Remove fractional part
|
||||
echo "Running benchmark with $NUM_PROMPTS prompts"
|
||||
|
||||
# Run the benchmark
|
||||
python "$SCRIPT_DIR/benchmark_serving_structured_output.py" $COMMON_PARAMS \
|
||||
--request-rate $qps \
|
||||
python "$SCRIPT_DIR/benchmark_serving_structured_output.py" "${COMMON_PARAMS[@]}" \
|
||||
--request-rate "$qps" \
|
||||
--result-filename "$FILENAME" \
|
||||
--num-prompts $NUM_PROMPTS
|
||||
--num-prompts "$NUM_PROMPTS"
|
||||
|
||||
echo "Completed benchmark with QPS: $qps"
|
||||
echo "----------------------------------------"
|
||||
|
||||
@@ -18,6 +18,7 @@ set(ENABLE_AVX512 $ENV{VLLM_CPU_AVX512})
|
||||
set(ENABLE_AVX512BF16 $ENV{VLLM_CPU_AVX512BF16})
|
||||
set(ENABLE_AVX512VNNI $ENV{VLLM_CPU_AVX512VNNI})
|
||||
set(ENABLE_AMXBF16 $ENV{VLLM_CPU_AMXBF16})
|
||||
set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
|
||||
|
||||
include_directories("${CMAKE_SOURCE_DIR}/csrc")
|
||||
|
||||
@@ -115,6 +116,10 @@ else()
|
||||
set(AVX512_FOUND ON)
|
||||
message(STATUS "AVX512 support enabled via VLLM_CPU_AVX512 environment variable")
|
||||
endif()
|
||||
if (ENABLE_ARM_BF16)
|
||||
set(ARM_BF16_FOUND ON)
|
||||
message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (AVX512_FOUND AND NOT AVX512_DISABLED)
|
||||
|
||||
@@ -19,7 +19,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
flashmla
|
||||
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
|
||||
GIT_TAG c2afa9cb93e674d5a9120a170a6da57b89267208
|
||||
GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
|
||||
+48
-23
@@ -2,33 +2,58 @@
|
||||
#include <torch/cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// This function assumes that `cpu_tensor` is a CPU tensor allocated with pinned
|
||||
// memory, and that UVA (Unified Virtual Addressing) is enabled.
|
||||
// This function assumes that `cpu_tensor` is a CPU tensor,
|
||||
// and that UVA (Unified Virtual Addressing) is enabled.
|
||||
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
|
||||
TORCH_CHECK(cpu_tensor.device().is_cpu(), "Input tensor must be on CPU");
|
||||
|
||||
// Get raw host pointer from CPU tensor
|
||||
void* host_ptr = cpu_tensor.data_ptr();
|
||||
// handle empty tensor
|
||||
if (cpu_tensor.numel() == 0) {
|
||||
return torch::empty(cpu_tensor.sizes(),
|
||||
cpu_tensor.options().device(torch::kCUDA));
|
||||
}
|
||||
|
||||
if (cpu_tensor.is_pinned()) {
|
||||
// If CPU tensor is pinned, directly get the device pointer.
|
||||
void* host_ptr = const_cast<void*>(cpu_tensor.data_ptr());
|
||||
void* device_ptr = nullptr;
|
||||
cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
|
||||
TORCH_CHECK(err == cudaSuccess,
|
||||
"cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
|
||||
|
||||
return torch::from_blob(
|
||||
device_ptr, cpu_tensor.sizes(), cpu_tensor.strides(),
|
||||
[base = cpu_tensor](void*) {}, // keep cpu tensor alive
|
||||
cpu_tensor.options().device(torch::kCUDA));
|
||||
}
|
||||
|
||||
// If CPU tensor is not pinned, allocate a new pinned memory buffer.
|
||||
torch::Tensor contiguous_cpu = cpu_tensor.contiguous();
|
||||
size_t nbytes = contiguous_cpu.nbytes();
|
||||
|
||||
void* host_ptr = nullptr;
|
||||
cudaError_t err = cudaHostAlloc(&host_ptr, nbytes, cudaHostAllocMapped);
|
||||
if (err != cudaSuccess) {
|
||||
AT_ERROR("cudaHostAlloc failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
err = cudaMemcpy(host_ptr, contiguous_cpu.data_ptr(), nbytes,
|
||||
cudaMemcpyDefault);
|
||||
if (err != cudaSuccess) {
|
||||
cudaFreeHost(host_ptr);
|
||||
AT_ERROR("cudaMemcpy failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
// Get a device pointer corresponding to the pinned host memory
|
||||
void* device_ptr = nullptr;
|
||||
cudaError_t err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
|
||||
TORCH_CHECK(err == cudaSuccess,
|
||||
"cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
|
||||
err = cudaHostGetDevicePointer(&device_ptr, host_ptr, 0);
|
||||
if (err != cudaSuccess) {
|
||||
cudaFreeHost(host_ptr);
|
||||
AT_ERROR("cudaHostGetDevicePointer failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
|
||||
// We'll use the same sizes, strides, and dtype as the CPU tensor.
|
||||
// TODO: check if layout is respected.
|
||||
auto sizes = cpu_tensor.sizes();
|
||||
auto strides = cpu_tensor.strides();
|
||||
auto options = cpu_tensor.options().device(torch::kCUDA);
|
||||
auto deleter = [host_ptr](void*) { cudaFreeHost(host_ptr); };
|
||||
|
||||
// use default no-op deleter, since the memory is owned by the original CPU
|
||||
// tensor
|
||||
torch::Tensor cuda_tensor =
|
||||
torch::from_blob(device_ptr, sizes, strides, options);
|
||||
|
||||
TORCH_CHECK(cuda_tensor.device().is_cuda(),
|
||||
"Resulting tensor is not on CUDA device");
|
||||
|
||||
return cuda_tensor;
|
||||
}
|
||||
return torch::from_blob(device_ptr, contiguous_cpu.sizes(),
|
||||
contiguous_cpu.strides(), deleter,
|
||||
contiguous_cpu.options().device(torch::kCUDA));
|
||||
}
|
||||
@@ -1568,8 +1568,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
{
|
||||
#endif
|
||||
unsigned int kOff = k + (thrd * A_CHUNK);
|
||||
unsigned int kOffcp =
|
||||
k_str + kOff; // min__(K - A_CHUNK, k_str + kOff);
|
||||
unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff);
|
||||
for (unsigned int n = 0; n < N; n += CHUNKK * sprdN) {
|
||||
__builtin_amdgcn_global_load_lds(
|
||||
(int*)(&A[min__(
|
||||
|
||||
+1
-1
@@ -349,7 +349,7 @@ void setup_kernel_smem_once() {
|
||||
void large_context_topk(
|
||||
const torch::Tensor& logits, torch::Tensor& indices,
|
||||
const torch::Tensor& seq_lens,
|
||||
c10::optional<torch::Tensor> row_starts = c10::nullopt) {
|
||||
std::optional<torch::Tensor> row_starts = std::nullopt) {
|
||||
TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
|
||||
TORCH_CHECK(indices.is_cuda(), "indices must be a CUDA tensor");
|
||||
TORCH_CHECK(seq_lens.is_cuda(), "seq_lens must be a CUDA tensor");
|
||||
|
||||
+15
-5
@@ -22,7 +22,12 @@
|
||||
# docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
|
||||
# =============================================================================
|
||||
|
||||
ARG CUDA_VERSION=12.9.1
|
||||
ARG CUDA_VERSION=12.8.1
|
||||
# BUILDER_CUDA_VERSION controls the CUDA toolkit used to compile csrc/ and
|
||||
# extensions (DeepGEMM, EP kernels). It can differ from CUDA_VERSION, which
|
||||
# is the CUDA version shipped in the final runtime image and used to select
|
||||
# the matching PyTorch wheel.
|
||||
ARG BUILDER_CUDA_VERSION=12.9.1
|
||||
ARG PYTHON_VERSION=3.12
|
||||
|
||||
# By parameterizing the base images, we allow third-party to use their own
|
||||
@@ -36,7 +41,7 @@ ARG PYTHON_VERSION=3.12
|
||||
# compatibility with other Linux OSes. The main reason for this is that the
|
||||
# glibc version is baked into the distro, and binaries built with one glibc
|
||||
# version are not backwards compatible with OSes that use an earlier version.
|
||||
ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
|
||||
ARG BUILD_BASE_IMAGE=nvidia/cuda:${BUILDER_CUDA_VERSION}-devel-ubuntu20.04
|
||||
# Using cuda base image with minimal dependencies necessary for JIT compilation (FlashInfer, DeepGEMM, EP kernels)
|
||||
ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-base-ubuntu22.04
|
||||
|
||||
@@ -92,8 +97,13 @@ ARG INSTALL_KV_CONNECTORS=false
|
||||
FROM ${BUILD_BASE_IMAGE} AS base
|
||||
|
||||
ARG CUDA_VERSION
|
||||
ARG BUILDER_CUDA_VERSION
|
||||
ARG PYTHON_VERSION
|
||||
|
||||
# Override the CUDA_VERSION env var inherited from the nvidia base image
|
||||
# (which equals BUILDER_CUDA_VERSION) so that $CUDA_VERSION in RUN commands
|
||||
# resolves to the runtime CUDA version used for PyTorch wheel selection.
|
||||
ENV CUDA_VERSION=${CUDA_VERSION}
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Install system dependencies including build tools
|
||||
@@ -133,7 +143,7 @@ ENV UV_LINK_MODE=copy
|
||||
RUN gcc --version
|
||||
|
||||
# Ensure CUDA compatibility library is loaded
|
||||
RUN echo "/usr/local/cuda-$(echo "$CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/cuda-compat.conf && ldconfig
|
||||
RUN echo "/usr/local/cuda-$(echo "$BUILDER_CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/cuda-compat.conf && ldconfig
|
||||
|
||||
# ============================================================
|
||||
# SLOW-CHANGING DEPENDENCIES BELOW
|
||||
@@ -309,7 +319,7 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
# Build DeepGEMM, pplx-kernels, DeepEP - runs in PARALLEL with csrc-build
|
||||
# This stage is independent and doesn't affect csrc cache
|
||||
FROM base AS extensions-build
|
||||
ARG CUDA_VERSION
|
||||
ARG BUILDER_CUDA_VERSION
|
||||
|
||||
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
@@ -325,7 +335,7 @@ COPY tools/install_deepgemm.sh /tmp/install_deepgemm.sh
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p /tmp/deepgemm/dist && \
|
||||
VLLM_DOCKER_BUILD_CONTEXT=1 TORCH_CUDA_ARCH_LIST="9.0a 10.0a" /tmp/install_deepgemm.sh \
|
||||
--cuda-version "${CUDA_VERSION}" \
|
||||
--cuda-version "${BUILDER_CUDA_VERSION}" \
|
||||
${DEEPGEMM_GIT_REF:+--ref "$DEEPGEMM_GIT_REF"} \
|
||||
--wheel-dir /tmp/deepgemm/dist || \
|
||||
echo "DeepGEMM build skipped (CUDA version requirement not met)"
|
||||
|
||||
@@ -20,6 +20,7 @@
|
||||
# VLLM_CPU_AVX512BF16=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AVX512VNNI=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AMXBF16=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_ARM_BF16=false (default)|true (for cross-compilation)
|
||||
#
|
||||
|
||||
######################### COMMON BASE IMAGE #########################
|
||||
@@ -108,9 +109,22 @@ ENV VLLM_CPU_AVX512VNNI=${VLLM_CPU_AVX512VNNI}
|
||||
# Support for building with AMXBF16 ISA: docker build --build-arg VLLM_CPU_AMXBF16="true" ...
|
||||
ARG VLLM_CPU_AMXBF16=1
|
||||
ENV VLLM_CPU_AMXBF16=${VLLM_CPU_AMXBF16}
|
||||
# Support for cross-compilation with ARM BF16 ISA: docker build --build-arg VLLM_CPU_ARM_BF16="true" ...
|
||||
ARG VLLM_CPU_ARM_BF16=0
|
||||
ENV VLLM_CPU_ARM_BF16=${VLLM_CPU_ARM_BF16}
|
||||
|
||||
WORKDIR /vllm-workspace
|
||||
|
||||
# Validate build arguments - prevent mixing incompatible ISA flags
|
||||
RUN if [ "$TARGETARCH" = "arm64" ] && { [ "$VLLM_CPU_AVX2" != "0" ] || [ "$VLLM_CPU_AVX512" != "0" ] || [ "$VLLM_CPU_AVX512BF16" != "0" ] || [ "$VLLM_CPU_AVX512VNNI" != "0" ]; }; then \
|
||||
echo "ERROR: Cannot use x86-specific ISA flags (AVX2, AVX512, etc.) when building for ARM64 (--platform=linux/arm64)"; \
|
||||
exit 1; \
|
||||
fi && \
|
||||
if [ "$TARGETARCH" = "amd64" ] && [ "$VLLM_CPU_ARM_BF16" != "0" ]; then \
|
||||
echo "ERROR: Cannot use ARM-specific ISA flags (ARM_BF16) when building for x86_64 (--platform=linux/amd64)"; \
|
||||
exit 1; \
|
||||
fi
|
||||
|
||||
# Copy build requirements
|
||||
COPY requirements/cpu-build.txt requirements/build.txt
|
||||
|
||||
@@ -224,6 +238,7 @@ ARG VLLM_CPU_AVX512
|
||||
ARG VLLM_CPU_AVX512BF16
|
||||
ARG VLLM_CPU_AVX512VNNI
|
||||
ARG VLLM_CPU_AMXBF16
|
||||
ARG VLLM_CPU_ARM_BF16
|
||||
ARG PYTHON_VERSION
|
||||
|
||||
LABEL ai.vllm.build.target-arch="${TARGETARCH}"
|
||||
@@ -233,6 +248,7 @@ LABEL ai.vllm.build.cpu-avx512="${VLLM_CPU_AVX512:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx512bf16="${VLLM_CPU_AVX512BF16:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx512vnni="${VLLM_CPU_AVX512VNNI:-false}"
|
||||
LABEL ai.vllm.build.cpu-amxbf16="${VLLM_CPU_AMXBF16:-false}"
|
||||
LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}"
|
||||
LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
@@ -2,6 +2,9 @@
|
||||
"_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
|
||||
"variable": {
|
||||
"CUDA_VERSION": {
|
||||
"default": "12.8.1"
|
||||
},
|
||||
"BUILDER_CUDA_VERSION": {
|
||||
"default": "12.9.1"
|
||||
},
|
||||
"PYTHON_VERSION": {
|
||||
@@ -11,7 +14,7 @@
|
||||
"default": "nvidia/cuda:12.9.1-devel-ubuntu20.04"
|
||||
},
|
||||
"FINAL_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:12.9.1-base-ubuntu22.04"
|
||||
"default": "nvidia/cuda:12.8.1-base-ubuntu22.04"
|
||||
},
|
||||
"GET_PIP_URL": {
|
||||
"default": "https://bootstrap.pypa.io/get-pip.py"
|
||||
|
||||
@@ -182,6 +182,7 @@ The following table lists backends that support full CUDA Graphs at the time of
|
||||
| FlashInfer | `UNIFORM_SINGLE_TOKEN_DECODE` | Will be set to `UNIFORM_BATCH` when using TRTLLM attention on Blackwell |
|
||||
| FlashMLA | `UNIFORM_BATCH` | |
|
||||
| FlashInferMLA | `UNIFORM_BATCH` | |
|
||||
| FlashInferMLASparse | `UNIFORM_BATCH` | |
|
||||
| AITER MLA | `UNIFORM_SINGLE_TOKEN_DECODE` | |
|
||||
| CUTLASS MLA | `UNIFORM_SINGLE_TOKEN_DECODE` | |
|
||||
| Mamba attention| `UNIFORM_SINGLE_TOKEN_DECODE` | |
|
||||
|
||||
@@ -109,6 +109,7 @@ Batch invariance has been tested and verified on the following models:
|
||||
- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
|
||||
- **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`
|
||||
- **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b`
|
||||
- **Mistral**: `mistralai/Mistral-7B-v0.3`
|
||||
|
||||
Other models may also work, but these have been explicitly validated. If you encounter issues with a specific model, please report them on the [GitHub issue tracker](https://github.com/vllm-project/vllm/issues/new/choose).
|
||||
|
||||
|
||||
@@ -84,7 +84,7 @@ Since simple RTN does not require data for weight quantization and the activatio
|
||||
Install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.9.2"
|
||||
pip install vllm "lm-eval[api]>=0.4.11"
|
||||
```
|
||||
|
||||
Load and run the model in `vllm`:
|
||||
|
||||
@@ -18,7 +18,7 @@ pip install llmcompressor
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.9.2"
|
||||
pip install vllm "lm-eval[api]>=0.4.11"
|
||||
```
|
||||
|
||||
## Quantization Process
|
||||
|
||||
@@ -23,7 +23,7 @@ pip install llmcompressor
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.9.2"
|
||||
pip install vllm "lm-eval[api]>=0.4.11"
|
||||
```
|
||||
|
||||
## Quantization Process
|
||||
|
||||
@@ -20,7 +20,7 @@ for more installation details.
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.9.2"
|
||||
pip install vllm "lm-eval[api]>=0.4.11"
|
||||
```
|
||||
|
||||
## Quantization Process
|
||||
|
||||
@@ -172,25 +172,78 @@ docker pull public.ecr.aws/q9t5s3a7/vllm-ci-postmerge-repo:${VLLM_COMMIT}-arm64-
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
# --8<-- [start:build-image-from-source]
|
||||
|
||||
## Building for your target ARM CPU
|
||||
|
||||
```bash
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--tag vllm-cpu-env .
|
||||
--platform=linux/arm64 \
|
||||
--build-arg VLLM_CPU_ARM_BF16=<false (default)|true> \
|
||||
--tag vllm-cpu-env \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
# Launching OpenAI server
|
||||
!!! note "Auto-detection by default"
|
||||
By default, ARM CPU instruction sets (BF16, NEON, etc.) are automatically detected from the build system's CPU flags. The `VLLM_CPU_ARM_BF16` build argument is used for cross-compilation:
|
||||
|
||||
- `VLLM_CPU_ARM_BF16=true` - Force-enable ARM BF16 support (build with BF16 regardless of build system capabilities)
|
||||
- `VLLM_CPU_ARM_BF16=false` - Rely on auto-detection (default)
|
||||
|
||||
### Examples
|
||||
|
||||
**Auto-detection build (native ARM)**
|
||||
|
||||
```bash
|
||||
# Building on ARM64 system - platform auto-detected
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--tag vllm-cpu-arm64 \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
**Cross-compile for ARM with BF16 support**
|
||||
|
||||
```bash
|
||||
# Building on ARM64 for newer ARM CPUs with BF16
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--build-arg VLLM_CPU_ARM_BF16=true \
|
||||
--tag vllm-cpu-arm64-bf16 \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
**Cross-compile from x86_64 to ARM64 with BF16**
|
||||
|
||||
```bash
|
||||
# Requires Docker buildx with ARM emulation (QEMU)
|
||||
docker buildx build -f docker/Dockerfile.cpu \
|
||||
--platform=linux/arm64 \
|
||||
--build-arg VLLM_CPU_ARM_BF16=true \
|
||||
--build-arg max_jobs=4 \
|
||||
--tag vllm-cpu-arm64-bf16 \
|
||||
--target vllm-openai \
|
||||
--load .
|
||||
```
|
||||
|
||||
!!! note "ARM BF16 requirements"
|
||||
ARM BF16 support requires ARMv8.6-A or later (FEAT_BF16). Supported on AWS Graviton3/4, AmpereOne, and other recent ARM processors.
|
||||
|
||||
## Launching the OpenAI server
|
||||
|
||||
```bash
|
||||
docker run --rm \
|
||||
--privileged=true \
|
||||
--security-opt seccomp=unconfined \
|
||||
--cap-add SYS_NICE \
|
||||
--shm-size=4g \
|
||||
-p 8000:8000 \
|
||||
-e VLLM_CPU_KVCACHE_SPACE=<KV cache space> \
|
||||
-e VLLM_CPU_OMP_THREADS_BIND=<CPU cores for inference> \
|
||||
vllm-cpu-env \
|
||||
--model=meta-llama/Llama-3.2-1B-Instruct \
|
||||
vllm-cpu-arm64 \
|
||||
meta-llama/Llama-3.2-1B-Instruct \
|
||||
--dtype=bfloat16 \
|
||||
other vLLM OpenAI server arguments
|
||||
```
|
||||
|
||||
!!! tip
|
||||
An alternative of `--privileged=true` is `--cap-add SYS_NICE --security-opt seccomp=unconfined`.
|
||||
!!! tip "Alternative to --privileged"
|
||||
Instead of `--privileged=true`, use `--cap-add SYS_NICE --security-opt seccomp=unconfined` for better security.
|
||||
|
||||
# --8<-- [end:build-image-from-source]
|
||||
# --8<-- [start:extra-information]
|
||||
|
||||
@@ -311,20 +311,31 @@ An OpenAI client example can be found here: [examples/pooling/embed/openai_embed
|
||||
|
||||
[ColBERT](https://arxiv.org/abs/2004.12832) (Contextualized Late Interaction over BERT) is a retrieval model that uses per-token embeddings and MaxSim scoring for document ranking. Unlike single-vector embedding models, ColBERT retains token-level representations and computes relevance scores through late interaction, providing better accuracy while being more efficient than cross-encoders.
|
||||
|
||||
vLLM supports ColBERT models for reranking tasks, automatically applying MaxSim scoring for query-document relevance:
|
||||
vLLM supports ColBERT models with multiple encoder backbones:
|
||||
|
||||
| Architecture | Backbone | Example HF Models |
|
||||
|---|---|---|
|
||||
| `HF_ColBERT` | BERT | `answerdotai/answerai-colbert-small-v1`, `colbert-ir/colbertv2.0` |
|
||||
| `ColBERTModernBertModel` | ModernBERT | `lightonai/GTE-ModernColBERT-v1` |
|
||||
| `ColBERTJinaRobertaModel` | Jina XLM-RoBERTa | `jinaai/jina-colbert-v2` |
|
||||
|
||||
**BERT-based ColBERT** models work out of the box:
|
||||
|
||||
```shell
|
||||
vllm serve answerdotai/answerai-colbert-small-v1
|
||||
```
|
||||
|
||||
Currently supports ColBERT models with standard BERT encoders (e.g., `answerdotai/answerai-colbert-small-v1`, `colbert-ir/colbertv2.0`).
|
||||
|
||||
ColBERT models with modified encoder architectures are not yet supported, including BERT variants with rotary embeddings (e.g., `jinaai/jina-colbert-v2`) or other custom encoders (e.g., `LiquidAI/LFM2-ColBERT-350M`).
|
||||
|
||||
If your standard BERT ColBERT model's config doesn't specify the architecture as `HF_ColBERT`, override it with:
|
||||
For **non-BERT backbones**, use `--hf-overrides` to set the correct architecture:
|
||||
|
||||
```shell
|
||||
vllm serve your-colbert-model --hf-overrides '{"architectures": ["HF_ColBERT"]}'
|
||||
# ModernBERT backbone
|
||||
vllm serve lightonai/GTE-ModernColBERT-v1 \
|
||||
--hf-overrides '{"architectures": ["ColBERTModernBertModel"]}'
|
||||
|
||||
# Jina XLM-RoBERTa backbone
|
||||
vllm serve jinaai/jina-colbert-v2 \
|
||||
--hf-overrides '{"architectures": ["ColBERTJinaRobertaModel"]}' \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
Then you can use the rerank endpoint:
|
||||
@@ -363,6 +374,77 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
|
||||
|
||||
An example can be found here: [examples/pooling/score/colbert_rerank_online.py](../../examples/pooling/score/colbert_rerank_online.py)
|
||||
|
||||
### ColQwen3 Multi-Modal Late Interaction Models
|
||||
|
||||
ColQwen3 is based on [ColPali](https://arxiv.org/abs/2407.01449), which extends ColBERT's late interaction approach to **multi-modal** inputs. While ColBERT operates on text-only token embeddings, ColPali/ColQwen3 can embed both **text and images** (e.g. PDF pages, screenshots, diagrams) into per-token L2-normalized vectors and compute relevance via MaxSim scoring. ColQwen3 specifically uses Qwen3-VL as its vision-language backbone.
|
||||
|
||||
| Architecture | Backbone | Example HF Models |
|
||||
|---|---|---|
|
||||
| `ColQwen3` | Qwen3-VL | `TomoroAI/tomoro-colqwen3-embed-4b`, `TomoroAI/tomoro-colqwen3-embed-8b` |
|
||||
| `OpsColQwen3Model` | Qwen3-VL | `OpenSearch-AI/Ops-Colqwen3-4B`, `OpenSearch-AI/Ops-Colqwen3-8B` |
|
||||
|
||||
Start the server:
|
||||
|
||||
```shell
|
||||
vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 4096
|
||||
```
|
||||
|
||||
Then you can use the rerank endpoint:
|
||||
|
||||
```shell
|
||||
curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
|
||||
"model": "TomoroAI/tomoro-colqwen3-embed-4b",
|
||||
"query": "What is machine learning?",
|
||||
"documents": [
|
||||
"Machine learning is a subset of artificial intelligence.",
|
||||
"Python is a programming language.",
|
||||
"Deep learning uses neural networks."
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
Or the score endpoint:
|
||||
|
||||
```shell
|
||||
curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
|
||||
"model": "TomoroAI/tomoro-colqwen3-embed-4b",
|
||||
"text_1": "What is the capital of France?",
|
||||
"text_2": ["The capital of France is Paris.", "Python is a programming language."]
|
||||
}'
|
||||
```
|
||||
|
||||
You can also get the raw token embeddings using the pooling endpoint with `token_embed` task:
|
||||
|
||||
```shell
|
||||
curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
|
||||
"model": "TomoroAI/tomoro-colqwen3-embed-4b",
|
||||
"input": "What is machine learning?",
|
||||
"task": "token_embed"
|
||||
}'
|
||||
```
|
||||
|
||||
For **image inputs**, use the chat-style `messages` field so that the vLLM multimodal processor handles them correctly:
|
||||
|
||||
```shell
|
||||
curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
|
||||
"model": "TomoroAI/tomoro-colqwen3-embed-4b",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
|
||||
{"type": "text", "text": "Describe the image."}
|
||||
]
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
Examples can be found here:
|
||||
|
||||
- Multi-vector retrieval: [examples/pooling/token_embed/colqwen3_token_embed_online.py](../../examples/pooling/token_embed/colqwen3_token_embed_online.py)
|
||||
- Reranking: [examples/pooling/score/colqwen3_rerank_online.py](../../examples/pooling/score/colqwen3_rerank_online.py)
|
||||
|
||||
### BAAI/bge-m3
|
||||
|
||||
The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json`
|
||||
|
||||
@@ -181,6 +181,7 @@ Some model architectures are supported via vLLM plugins. These plugins extend vL
|
||||
| Architecture | Models | Plugin Repository |
|
||||
|--------------|--------|-------------------|
|
||||
| `BartForConditionalGeneration` | BART | [bart-plugin](https://github.com/vllm-project/bart-plugin) |
|
||||
| `Florence2ForConditionalGeneration` | Florence-2 | [bart-plugin](https://github.com/vllm-project/bart-plugin) |
|
||||
|
||||
For other model architectures not natively supported, in particular for Encoder-Decoder models, we recommend following a similar pattern by implementing support through the plugin system.
|
||||
|
||||
@@ -728,6 +729,8 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `OpenPanguVLForConditionalGeneration` | openpangu-VL | T + I<sup>E+</sup> + V<sup>E+</sup> |`FreedomIntelligence/openPangu-VL-7B` | ✅︎ | ✅︎ |
|
||||
| `Ovis` | Ovis2, Ovis1.6 | T + I<sup>+</sup> | `AIDC-AI/Ovis2-1B`, `AIDC-AI/Ovis1.6-Llama3.2-3B`, etc. | | ✅︎ |
|
||||
| `Ovis2_5` | Ovis2.5 | T + I<sup>+</sup> + V | `AIDC-AI/Ovis2.5-9B`, etc. | | |
|
||||
| `Ovis2_6ForCausalLM` | Ovis2.6 | T + I<sup>+</sup> + V | `AIDC-AI/Ovis2.6-2B`, etc. | | |
|
||||
| `Ovis2_6_MoeForCausalLM` | Ovis2.6 | T + I<sup>+</sup> + V | `AIDC-AI/Ovis2.6-30B-A3B`, etc. | | |
|
||||
| `PaddleOCRVLForConditionalGeneration` | Paddle-OCR | T + I<sup>+</sup> | `PaddlePaddle/PaddleOCR-VL`, etc. | | |
|
||||
| `PaliGemmaForConditionalGeneration` | PaliGemma, PaliGemma 2 | T + I<sup>E</sup> | `google/paligemma-3b-pt-224`, `google/paligemma-3b-mix-224`, `google/paligemma2-3b-ft-docci-448`, etc. | ✅︎ | ✅︎ |
|
||||
| `Phi3VForCausalLM` | Phi-3-Vision, Phi-3.5-Vision | T + I<sup>E+</sup> | `microsoft/Phi-3-vision-128k-instruct`, `microsoft/Phi-3.5-vision-instruct`, etc. | | ✅︎ |
|
||||
|
||||
@@ -137,6 +137,7 @@ Please note that prefix caching is not yet supported for any of the above models
|
||||
Whisper is supported natively. Other encoder-decoder models are supported via the plugin system:
|
||||
|
||||
- **BART**: `BartForConditionalGeneration` is supported via the official [bart-plugin](https://github.com/vllm-project/bart-plugin).
|
||||
- **Florence-2**: `Florence2ForConditionalGeneration` is supported via the official [bart-plugin](https://github.com/vllm-project/bart-plugin).
|
||||
|
||||
For other encoder-decoder models (e.g., `MllamaForConditionalGeneration`), we recommend
|
||||
following a similar pattern by implementing support through the [plugin system](../design/plugin_system.md).
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Test pause/resume with Data Parallel (DP) via HTTP API.
|
||||
|
||||
This example demonstrates coordinated pause/resume across multiple DP ranks.
|
||||
The pause synchronizes across all DP engines via all-reduce.
|
||||
|
||||
Prerequisites:
|
||||
Start a vLLM server with data parallelism:
|
||||
|
||||
$ VLLM_SERVER_DEV_MODE=1 vllm serve facebook/opt-125m \
|
||||
--enforce-eager \
|
||||
--data-parallel-size 4 \
|
||||
--tensor-parallel-size 1
|
||||
|
||||
Then run this script:
|
||||
|
||||
$ python data_parallel_pause_resume.py
|
||||
|
||||
The test verifies pause works by:
|
||||
1. Starting a streaming generation request
|
||||
2. Pausing the server mid-generation
|
||||
3. Sleeping for PAUSE_DURATION seconds
|
||||
4. Resuming the server
|
||||
5. Verifying there was a gap in token generation matching the pause duration
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import threading
|
||||
import time
|
||||
|
||||
import requests
|
||||
from openai import OpenAI
|
||||
|
||||
BASE_URL = "http://localhost:8000"
|
||||
MODEL_NAME = "facebook/opt-125m"
|
||||
PAUSE_DURATION = 3.0
|
||||
|
||||
|
||||
def pause_generation(base_url: str, mode: str = "keep") -> None:
|
||||
"""Pause generation via HTTP endpoint."""
|
||||
url = f"{base_url}/pause"
|
||||
response = requests.post(url, params={"mode": mode}, timeout=60)
|
||||
response.raise_for_status()
|
||||
print("Server paused")
|
||||
|
||||
|
||||
def resume_generation(base_url: str) -> None:
|
||||
"""Resume generation via HTTP endpoint."""
|
||||
url = f"{base_url}/resume"
|
||||
response = requests.post(url, timeout=60)
|
||||
response.raise_for_status()
|
||||
print("Server resumed")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--base-url", default=BASE_URL)
|
||||
parser.add_argument("--model", default=MODEL_NAME)
|
||||
args = parser.parse_args()
|
||||
|
||||
client = OpenAI(
|
||||
base_url=f"{args.base_url}/v1",
|
||||
api_key="EMPTY",
|
||||
)
|
||||
|
||||
prompt = "Write a long story about a dragon. Once upon a time"
|
||||
token_times: list[float] = []
|
||||
pause_token_idx = 0
|
||||
pause_triggered = threading.Event()
|
||||
|
||||
def generator_thread():
|
||||
"""Stream tokens and record timestamps."""
|
||||
stream = client.completions.create(
|
||||
model=args.model,
|
||||
prompt=prompt,
|
||||
max_tokens=50,
|
||||
stream=True,
|
||||
)
|
||||
for chunk in stream:
|
||||
if chunk.choices[0].text:
|
||||
token_times.append(time.monotonic())
|
||||
token_count = len(token_times)
|
||||
print(f"Token {token_count}: {chunk.choices[0].text!r}")
|
||||
|
||||
# Signal controller after some tokens
|
||||
if token_count >= 5 and not pause_triggered.is_set():
|
||||
pause_triggered.set()
|
||||
|
||||
def controller_thread():
|
||||
"""Pause and resume the server."""
|
||||
nonlocal pause_token_idx
|
||||
|
||||
# Wait for some tokens
|
||||
pause_triggered.wait()
|
||||
|
||||
print(f"\nPausing server (keep mode) at token {len(token_times)}...")
|
||||
pause_generation(args.base_url, mode="keep")
|
||||
pause_token_idx = len(token_times)
|
||||
print(f"Sleeping for {PAUSE_DURATION}s...")
|
||||
|
||||
time.sleep(PAUSE_DURATION)
|
||||
|
||||
print("Resuming server...")
|
||||
resume_generation(args.base_url)
|
||||
print("Resumed!\n")
|
||||
|
||||
# Run both threads
|
||||
gen_thread = threading.Thread(target=generator_thread)
|
||||
ctrl_thread = threading.Thread(target=controller_thread)
|
||||
|
||||
gen_thread.start()
|
||||
ctrl_thread.start()
|
||||
|
||||
gen_thread.join()
|
||||
ctrl_thread.join()
|
||||
|
||||
# Check gap at the pause point
|
||||
if pause_token_idx < len(token_times):
|
||||
pause_gap = token_times[pause_token_idx] - token_times[pause_token_idx - 1]
|
||||
print(
|
||||
f"\nGap after pause (token {pause_token_idx} -> "
|
||||
f"{pause_token_idx + 1}): {pause_gap:.3f}s"
|
||||
)
|
||||
if pause_gap >= PAUSE_DURATION * 0.9:
|
||||
print("Test passed! Pause synchronized across DP ranks.")
|
||||
else:
|
||||
print(f"Test failed! Expected ~{PAUSE_DURATION}s gap, got {pause_gap:.3f}s")
|
||||
else:
|
||||
print("Test failed! No tokens were generated after resuming.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -8,7 +8,7 @@ declare -a PIDS=()
|
||||
###############################################################################
|
||||
MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
|
||||
LOG_PATH="${LOG_PATH:-./logs}"
|
||||
mkdir -p $LOG_PATH
|
||||
mkdir -p "$LOG_PATH"
|
||||
|
||||
ENCODE_PORT="${ENCODE_PORT:-19534}"
|
||||
PREFILL_PORT="${PREFILL_PORT:-19535}"
|
||||
@@ -84,10 +84,10 @@ trap cleanup TERM
|
||||
|
||||
# clear previous cache
|
||||
echo "remove previous ec cache folder"
|
||||
rm -rf $EC_SHARED_STORAGE_PATH
|
||||
rm -rf "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
echo "make ec cache folder"
|
||||
mkdir -p $EC_SHARED_STORAGE_PATH
|
||||
mkdir -p "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
###############################################################################
|
||||
# Encoder worker
|
||||
@@ -100,7 +100,7 @@ CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
|
||||
--no-enable-prefix-caching \
|
||||
--max-num-batched-tokens 114688 \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_producer",
|
||||
@@ -124,7 +124,7 @@ vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--enable-request-id-headers \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_consumer",
|
||||
@@ -152,7 +152,7 @@ vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--enable-request-id-headers \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_consumer"
|
||||
@@ -162,9 +162,9 @@ vllm serve "$MODEL" \
|
||||
PIDS+=($!)
|
||||
|
||||
# Wait for workers
|
||||
wait_for_server $ENCODE_PORT
|
||||
wait_for_server $PREFILL_PORT
|
||||
wait_for_server $DECODE_PORT
|
||||
wait_for_server "$ENCODE_PORT"
|
||||
wait_for_server "$PREFILL_PORT"
|
||||
wait_for_server "$DECODE_PORT"
|
||||
|
||||
###############################################################################
|
||||
# Proxy
|
||||
@@ -179,7 +179,7 @@ python disagg_epd_proxy.py \
|
||||
|
||||
PIDS+=($!)
|
||||
|
||||
wait_for_server $PROXY_PORT
|
||||
wait_for_server "$PROXY_PORT"
|
||||
echo "All services are up!"
|
||||
|
||||
###############################################################################
|
||||
@@ -187,14 +187,14 @@ echo "All services are up!"
|
||||
###############################################################################
|
||||
echo "Running benchmark (stream)..."
|
||||
vllm bench serve \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--backend openai-chat \
|
||||
--endpoint /v1/chat/completions \
|
||||
--dataset-name hf \
|
||||
--dataset-path lmarena-ai/VisionArena-Chat \
|
||||
--seed 0 \
|
||||
--num-prompts $NUM_PROMPTS \
|
||||
--port $PROXY_PORT
|
||||
--num-prompts "$NUM_PROMPTS" \
|
||||
--port "$PROXY_PORT"
|
||||
|
||||
PIDS+=($!)
|
||||
|
||||
@@ -202,10 +202,10 @@ PIDS+=($!)
|
||||
# Single request with local image
|
||||
###############################################################################
|
||||
echo "Running single request with local image (non-stream)..."
|
||||
curl http://127.0.0.1:${PROXY_PORT}/v1/chat/completions \
|
||||
curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "'${MODEL}'",
|
||||
"model": "'"${MODEL}"'",
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": [
|
||||
|
||||
@@ -8,7 +8,7 @@ declare -a PIDS=()
|
||||
###############################################################################
|
||||
MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
|
||||
LOG_PATH="${LOG_PATH:-./logs}"
|
||||
mkdir -p $LOG_PATH
|
||||
mkdir -p "$LOG_PATH"
|
||||
|
||||
ENCODE_PORT="${ENCODE_PORT:-19534}"
|
||||
PREFILL_DECODE_PORT="${PREFILL_DECODE_PORT:-19535}"
|
||||
@@ -78,10 +78,10 @@ trap cleanup TERM
|
||||
|
||||
# clear previous cache
|
||||
echo "remove previous ec cache folder"
|
||||
rm -rf $EC_SHARED_STORAGE_PATH
|
||||
rm -rf "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
echo "make ec cache folder"
|
||||
mkdir -p $EC_SHARED_STORAGE_PATH
|
||||
mkdir -p "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
###############################################################################
|
||||
# Encoder worker
|
||||
@@ -94,7 +94,7 @@ CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
|
||||
--no-enable-prefix-caching \
|
||||
--max-num-batched-tokens 114688 \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_producer",
|
||||
@@ -115,7 +115,7 @@ CUDA_VISIBLE_DEVICES="$GPU_PD" vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--enable-request-id-headers \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_consumer",
|
||||
@@ -128,8 +128,8 @@ CUDA_VISIBLE_DEVICES="$GPU_PD" vllm serve "$MODEL" \
|
||||
PIDS+=($!)
|
||||
|
||||
# Wait for workers
|
||||
wait_for_server $ENCODE_PORT
|
||||
wait_for_server $PREFILL_DECODE_PORT
|
||||
wait_for_server "$ENCODE_PORT"
|
||||
wait_for_server "$PREFILL_DECODE_PORT"
|
||||
|
||||
###############################################################################
|
||||
# Proxy
|
||||
@@ -144,7 +144,7 @@ python disagg_epd_proxy.py \
|
||||
|
||||
PIDS+=($!)
|
||||
|
||||
wait_for_server $PROXY_PORT
|
||||
wait_for_server "$PROXY_PORT"
|
||||
echo "All services are up!"
|
||||
|
||||
###############################################################################
|
||||
@@ -152,14 +152,14 @@ echo "All services are up!"
|
||||
###############################################################################
|
||||
echo "Running benchmark (stream)..."
|
||||
vllm bench serve \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--backend openai-chat \
|
||||
--endpoint /v1/chat/completions \
|
||||
--dataset-name hf \
|
||||
--dataset-path lmarena-ai/VisionArena-Chat \
|
||||
--seed 0 \
|
||||
--num-prompts $NUM_PROMPTS \
|
||||
--port $PROXY_PORT
|
||||
--num-prompts "$NUM_PROMPTS" \
|
||||
--port "$PROXY_PORT"
|
||||
|
||||
PIDS+=($!)
|
||||
|
||||
@@ -167,10 +167,10 @@ PIDS+=($!)
|
||||
# Single request with local image
|
||||
###############################################################################
|
||||
echo "Running single request with local image (non-stream)..."
|
||||
curl http://127.0.0.1:${PROXY_PORT}/v1/chat/completions \
|
||||
curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "'${MODEL}'",
|
||||
"model": "'"${MODEL}"'",
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": [
|
||||
|
||||
@@ -54,7 +54,7 @@ wait_for_server() {
|
||||
# You can also adjust --kv-ip and --kv-port for distributed inference.
|
||||
|
||||
# prefilling instance, which is the KV producer
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve $MODEL_NAME \
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve "$MODEL_NAME" \
|
||||
--host 0.0.0.0 \
|
||||
--port 8100 \
|
||||
--max-model-len 100 \
|
||||
@@ -64,7 +64,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve $MODEL_NAME \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":"1e9","kv_port":"14579","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8100","send_type":"PUT_ASYNC"}}' &
|
||||
|
||||
# decoding instance, which is the KV consumer
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve $MODEL_NAME \
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve "$MODEL_NAME" \
|
||||
--host 0.0.0.0 \
|
||||
--port 8200 \
|
||||
--max-model-len 100 \
|
||||
|
||||
@@ -328,9 +328,9 @@ class Proxy:
|
||||
if instance_type == "decode" and instance in self.decode_instances:
|
||||
self.decode_instances.remove(instance)
|
||||
self.decode_cycler = itertools.cycle(self.decode_instances)
|
||||
if instance_type == "prefill" and instance in self.decode_instances:
|
||||
if instance_type == "prefill" and instance in self.prefill_instances:
|
||||
self.prefill_instances.remove(instance)
|
||||
self.prefill_cycler = itertools.cycle(self.decode_instances)
|
||||
self.prefill_cycler = itertools.cycle(self.prefill_instances)
|
||||
|
||||
|
||||
class RoundRobinSchedulingPolicy(SchedulingPolicy):
|
||||
|
||||
@@ -34,7 +34,7 @@ wait_for_server() {
|
||||
done" && return 0 || return 1
|
||||
}
|
||||
|
||||
vllm serve $MODEL_NAME \
|
||||
vllm serve "$MODEL_NAME" \
|
||||
--port 8100 \
|
||||
--max-model-len 100 \
|
||||
--enforce-eager \
|
||||
|
||||
+12
-11
@@ -143,7 +143,7 @@ main() {
|
||||
IFS=',' read -ra BOOTSTRAP_PORT_ARRAY <<< "$BOOTSTRAP_PORTS"
|
||||
IFS=',' read -ra DECODE_PORT_ARRAY <<< "$DECODE_PORTS"
|
||||
|
||||
proxy_param=""
|
||||
proxy_args=()
|
||||
|
||||
# =============================================================================
|
||||
# Launch Prefill Servers (X Producers)
|
||||
@@ -156,12 +156,12 @@ main() {
|
||||
local bootstrap_port=${BOOTSTRAP_PORT_ARRAY[$i]}
|
||||
|
||||
echo " Prefill server $((i+1)): GPU $gpu_id, Port $port, Bootstrap Port $bootstrap_port"
|
||||
VLLM_MOONCAKE_BOOTSTRAP_PORT=$bootstrap_port CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \
|
||||
--port $port \
|
||||
VLLM_MOONCAKE_BOOTSTRAP_PORT=$bootstrap_port CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
|
||||
--port "$port" \
|
||||
--kv-transfer-config \
|
||||
"{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_producer\"}" > prefill$((i+1)).log 2>&1 &
|
||||
PIDS+=($!)
|
||||
proxy_param="${proxy_param} --prefill http://0.0.0.0:${port} $bootstrap_port"
|
||||
proxy_args+=(--prefill "http://0.0.0.0:${port}" "$bootstrap_port")
|
||||
done
|
||||
|
||||
# =============================================================================
|
||||
@@ -174,12 +174,12 @@ main() {
|
||||
local port=${DECODE_PORT_ARRAY[$i]}
|
||||
|
||||
echo " Decode server $((i+1)): GPU $gpu_id, Port $port"
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \
|
||||
--port $port \
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
|
||||
--port "$port" \
|
||||
--kv-transfer-config \
|
||||
"{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_consumer\"}" > decode$((i+1)).log 2>&1 &
|
||||
PIDS+=($!)
|
||||
proxy_param="${proxy_param} --decode http://0.0.0.0:${port}"
|
||||
proxy_args+=(--decode "http://0.0.0.0:${port}")
|
||||
done
|
||||
|
||||
# =============================================================================
|
||||
@@ -187,7 +187,7 @@ main() {
|
||||
# =============================================================================
|
||||
echo ""
|
||||
echo "Starting proxy server on port $PROXY_PORT..."
|
||||
python3 mooncake_connector_proxy.py $proxy_param --port $PROXY_PORT > proxy.log 2>&1 &
|
||||
python3 mooncake_connector_proxy.py "${proxy_args[@]}" --port "$PROXY_PORT" > proxy.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# =============================================================================
|
||||
@@ -196,9 +196,10 @@ main() {
|
||||
echo ""
|
||||
echo "Waiting for all servers to start..."
|
||||
for port in "${PREFILL_PORT_ARRAY[@]}" "${DECODE_PORT_ARRAY[@]}"; do
|
||||
if ! wait_for_server $port; then
|
||||
if ! wait_for_server "$port"; then
|
||||
echo "Failed to start server on port $port"
|
||||
cleanup
|
||||
# shellcheck disable=SC2317
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
@@ -209,8 +210,8 @@ main() {
|
||||
# =============================================================================
|
||||
# Run Benchmark
|
||||
# =============================================================================
|
||||
vllm bench serve --port $PROXY_PORT --seed $(date +%s) \
|
||||
--backend vllm --model $MODEL \
|
||||
vllm bench serve --port "$PROXY_PORT" --seed "$(date +%s)" \
|
||||
--backend vllm --model "$MODEL" \
|
||||
--dataset-name random --random-input-len 7500 --random-output-len 200 \
|
||||
--num-prompts 200 --burstiness 100 --request-rate 2 | tee benchmark.log
|
||||
|
||||
|
||||
+8
-7
@@ -166,10 +166,10 @@ main() {
|
||||
local kv_port=$((21001 + i))
|
||||
|
||||
echo " Prefill server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--host 0.0.0.0 \
|
||||
--port $port \
|
||||
--port "$port" \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--dtype float16 \
|
||||
@@ -194,10 +194,10 @@ main() {
|
||||
local kv_port=$((22001 + i))
|
||||
|
||||
echo " Decode server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--host 0.0.0.0 \
|
||||
--port $port \
|
||||
--port "$port" \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--dtype float16 \
|
||||
@@ -217,9 +217,10 @@ main() {
|
||||
echo ""
|
||||
echo "Waiting for all servers to start..."
|
||||
for port in "${PREFILL_PORT_ARRAY[@]}" "${DECODE_PORT_ARRAY[@]}"; do
|
||||
if ! wait_for_server $port; then
|
||||
if ! wait_for_server "$port"; then
|
||||
echo "Failed to start server on port $port"
|
||||
cleanup
|
||||
# shellcheck disable=SC2317
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
@@ -231,8 +232,8 @@ main() {
|
||||
# Run Benchmark
|
||||
# =============================================================================
|
||||
cd ../../../benchmarks/
|
||||
vllm bench serve --port 10001 --seed $(date +%s) \
|
||||
--model $MODEL \
|
||||
vllm bench serve --port 10001 --seed "$(date +%s)" \
|
||||
--model "$MODEL" \
|
||||
--dataset-name random --random-input-len 7500 --random-output-len 200 \
|
||||
--num-prompts 200 --burstiness 100 --request-rate 2 | tee benchmark.log
|
||||
|
||||
|
||||
@@ -50,8 +50,8 @@ while [[ $# -gt 0 ]]; do
|
||||
done
|
||||
|
||||
vllm bench serve \
|
||||
--model $MODEL_NAME \
|
||||
--host $HOST \
|
||||
--port $PORT \
|
||||
--num-prompts $NUM_PROMPTS \
|
||||
--request-rate $REQUEST_RATE
|
||||
--model "$MODEL_NAME" \
|
||||
--host "$HOST" \
|
||||
--port "$PORT" \
|
||||
--num-prompts "$NUM_PROMPTS" \
|
||||
--request-rate "$REQUEST_RATE"
|
||||
|
||||
@@ -57,15 +57,15 @@ echo "Starting vLLM server for $MODEL_NAME with data parallel size: $DATA_PARALL
|
||||
export RAY_DEDUP_LOGS=0
|
||||
export VLLM_USE_DEEP_GEMM=1
|
||||
|
||||
vllm serve $MODEL_NAME \
|
||||
--data-parallel-size $DATA_PARALLEL_SIZE \
|
||||
--data-parallel-size-local $DATA_PARALLEL_SIZE \
|
||||
vllm serve "$MODEL_NAME" \
|
||||
--data-parallel-size "$DATA_PARALLEL_SIZE" \
|
||||
--data-parallel-size-local "$DATA_PARALLEL_SIZE" \
|
||||
--data-parallel-backend ray \
|
||||
--enforce-eager \
|
||||
--enable-expert-parallel \
|
||||
--enable-eplb \
|
||||
--all2all-backend pplx \
|
||||
--num-redundant-experts $REDUNDANT_EXPERTS \
|
||||
--num-redundant-experts "$REDUNDANT_EXPERTS" \
|
||||
--trust-remote-code \
|
||||
--host $HOST \
|
||||
--port $PORT
|
||||
--host "$HOST" \
|
||||
--port "$PORT"
|
||||
|
||||
@@ -57,8 +57,7 @@ case "$subcommand" in
|
||||
|
||||
# Retry until the worker node connects to the head node or the timeout expires.
|
||||
for (( i=0; i < $ray_init_timeout; i+=5 )); do
|
||||
ray start --address=$ray_address:$ray_port --block "${start_params[@]}"
|
||||
if [ $? -eq 0 ]; then
|
||||
if ray start --address="$ray_address":"$ray_port" --block "${start_params[@]}"; then
|
||||
echo "Worker: Ray runtime started with head address $ray_address:$ray_port"
|
||||
exit 0
|
||||
fi
|
||||
@@ -95,12 +94,12 @@ case "$subcommand" in
|
||||
fi
|
||||
|
||||
# Start the Ray head node.
|
||||
ray start --head --port=$ray_port "${start_params[@]}"
|
||||
ray start --head --port="$ray_port" "${start_params[@]}"
|
||||
|
||||
# Poll Ray until every worker node is active.
|
||||
for (( i=0; i < $ray_init_timeout; i+=5 )); do
|
||||
active_nodes=`python3 -c 'import ray; ray.init(); print(sum(node["Alive"] for node in ray.nodes()))'`
|
||||
if [ $active_nodes -eq $ray_cluster_size ]; then
|
||||
active_nodes=$(python3 -c 'import ray; ray.init(); print(sum(node["Alive"] for node in ray.nodes()))')
|
||||
if [ "$active_nodes" -eq "$ray_cluster_size" ]; then
|
||||
echo "All ray workers are active and the ray cluster is initialized successfully."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
@@ -22,11 +22,10 @@ check_hf_token() {
|
||||
|
||||
check_num_gpus() {
|
||||
# can you check if the number of GPUs are >=2 via nvidia-smi/rocm-smi?
|
||||
which rocm-smi > /dev/null 2>&1
|
||||
if [ $? -ne 0 ]; then
|
||||
if ! which rocm-smi > /dev/null 2>&1; then
|
||||
num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
|
||||
else
|
||||
num_gpus=$(rocm-smi --showid | grep Instinct | wc -l)
|
||||
num_gpus=$(rocm-smi --showid | grep -c Instinct)
|
||||
fi
|
||||
|
||||
if [ "$num_gpus" -lt 2 ]; then
|
||||
@@ -39,8 +38,7 @@ check_num_gpus() {
|
||||
|
||||
ensure_python_library_installed() {
|
||||
echo "Checking if $1 is installed..."
|
||||
python3 -c "import $1" > /dev/null 2>&1
|
||||
if [ $? -ne 0 ]; then
|
||||
if ! python3 -c "import $1" > /dev/null 2>&1; then
|
||||
if [ "$1" == "nixl" ]; then
|
||||
echo "$1 is not installed. Please refer to https://github.com/ai-dynamo/nixl for installation."
|
||||
else
|
||||
@@ -102,12 +100,12 @@ main() {
|
||||
bash disagg_vllm_launcher.sh prefiller \
|
||||
> >(tee prefiller.log) 2>&1 &
|
||||
prefiller_pid=$!
|
||||
PIDS+=($prefiller_pid)
|
||||
PIDS+=("$prefiller_pid")
|
||||
|
||||
bash disagg_vllm_launcher.sh decoder \
|
||||
> >(tee decoder.log) 2>&1 &
|
||||
decoder_pid=$!
|
||||
PIDS+=($decoder_pid)
|
||||
PIDS+=("$decoder_pid")
|
||||
|
||||
python3 disagg_proxy_server.py \
|
||||
--host localhost \
|
||||
@@ -118,7 +116,7 @@ main() {
|
||||
--decoder-port 8200 \
|
||||
> >(tee proxy.log) 2>&1 &
|
||||
proxy_pid=$!
|
||||
PIDS+=($proxy_pid)
|
||||
PIDS+=("$proxy_pid")
|
||||
|
||||
wait_for_server 8100
|
||||
wait_for_server 8200
|
||||
@@ -128,7 +126,7 @@ main() {
|
||||
|
||||
# begin benchmark
|
||||
cd ../../../../benchmarks/
|
||||
vllm bench serve --port 9000 --seed $(date +%s) \
|
||||
vllm bench serve --port 9000 --seed "$(date +%s)" \
|
||||
--model meta-llama/Llama-3.1-8B-Instruct \
|
||||
--dataset-name random --random-input-len 7500 --random-output-len 200 \
|
||||
--num-prompts 200 --burstiness 100 --request-rate 3.6 | tee benchmark.log
|
||||
|
||||
@@ -34,7 +34,7 @@ if [[ $1 == "prefiller" ]]; then
|
||||
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
CUDA_VISIBLE_DEVICES=0 \
|
||||
vllm serve $MODEL \
|
||||
vllm serve "$MODEL" \
|
||||
--port 8100 \
|
||||
--enforce-eager \
|
||||
--kv-transfer-config \
|
||||
@@ -51,7 +51,7 @@ elif [[ $1 == "decoder" ]]; then
|
||||
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
CUDA_VISIBLE_DEVICES=1 \
|
||||
vllm serve $MODEL \
|
||||
vllm serve "$MODEL" \
|
||||
--port 8200 \
|
||||
--enforce-eager \
|
||||
--kv-transfer-config \
|
||||
|
||||
@@ -103,7 +103,7 @@ vllm serve "$MODEL_NAME" \
|
||||
--tensor-parallel-size "$GPU_COUNT" \
|
||||
--enforce-eager \
|
||||
--pooler-config "$POOLER_CONFIG" \
|
||||
--served-model-name ${MODEL_CODE} \
|
||||
--served-model-name "${MODEL_CODE}" \
|
||||
--api-key "$API_KEY" \
|
||||
--trust-remote-code \
|
||||
--port "$PORT" \
|
||||
|
||||
@@ -20,15 +20,17 @@ def main():
|
||||
torch.set_default_dtype(torch.float16)
|
||||
image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" # noqa: E501
|
||||
|
||||
img_prompt = dict(
|
||||
img_data = dict(
|
||||
data=image_url,
|
||||
data_format="url",
|
||||
image_format="tiff",
|
||||
out_data_format="b64_json",
|
||||
)
|
||||
|
||||
prompt = dict(data=img_data)
|
||||
|
||||
llm = LLM(
|
||||
model="christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
|
||||
model="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
|
||||
skip_tokenizer_init=True,
|
||||
trust_remote_code=True,
|
||||
enforce_eager=True,
|
||||
@@ -41,7 +43,7 @@ def main():
|
||||
enable_mm_embeds=True,
|
||||
)
|
||||
|
||||
pooler_output = llm.encode(img_prompt, pooling_task="plugin")
|
||||
pooler_output = llm.encode(prompt, pooling_task="plugin")
|
||||
output = pooler_output[0].outputs
|
||||
|
||||
print(output)
|
||||
|
||||
@@ -391,7 +391,7 @@ if __name__ == "__main__":
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
|
||||
default="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
|
||||
help="Path to a checkpoint file to load from.",
|
||||
)
|
||||
parser.add_argument(
|
||||
|
||||
@@ -14,9 +14,7 @@ import requests
|
||||
# - install TerraTorch v1.1 (or later):
|
||||
# pip install terratorch>=v1.1
|
||||
# - start vllm in serving mode with the below args
|
||||
# --model='christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM'
|
||||
# --model-impl terratorch
|
||||
# --trust-remote-code
|
||||
# --model='ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11'
|
||||
# --skip-tokenizer-init --enforce-eager
|
||||
# --io-processor-plugin terratorch_segmentation
|
||||
# --enable-mm-embeds
|
||||
@@ -34,7 +32,7 @@ def main():
|
||||
"out_data_format": "b64_json",
|
||||
},
|
||||
"priority": 0,
|
||||
"model": "christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
|
||||
"model": "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
|
||||
}
|
||||
|
||||
ret = requests.post(server_endpoint, json=request_payload_url)
|
||||
|
||||
@@ -1,15 +1,27 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Example of using ColBERT late interaction model for reranking.
|
||||
Example of using ColBERT late interaction models for reranking and scoring.
|
||||
|
||||
ColBERT (Contextualized Late Interaction over BERT) uses per-token embeddings
|
||||
and MaxSim scoring for document reranking, providing better accuracy than
|
||||
single-vector models while being more efficient than cross-encoders.
|
||||
|
||||
Start the server with:
|
||||
vLLM supports ColBERT with multiple encoder backbones. Start the server
|
||||
with one of the following:
|
||||
|
||||
# BERT backbone (works out of the box)
|
||||
vllm serve answerdotai/answerai-colbert-small-v1
|
||||
|
||||
# ModernBERT backbone
|
||||
vllm serve lightonai/GTE-ModernColBERT-v1 \
|
||||
--hf-overrides '{"architectures": ["ColBERTModernBertModel"]}'
|
||||
|
||||
# Jina XLM-RoBERTa backbone
|
||||
vllm serve jinaai/jina-colbert-v2 \
|
||||
--hf-overrides '{"architectures": ["ColBERTJinaRobertaModel"]}' \
|
||||
--trust-remote-code
|
||||
|
||||
Then run this script:
|
||||
python colbert_rerank_online.py
|
||||
"""
|
||||
@@ -18,39 +30,62 @@ import json
|
||||
|
||||
import requests
|
||||
|
||||
url = "http://127.0.0.1:8000/rerank"
|
||||
# Change this to match the model you started the server with
|
||||
MODEL = "answerdotai/answerai-colbert-small-v1"
|
||||
BASE_URL = "http://127.0.0.1:8000"
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
data = {
|
||||
"model": "answerdotai/answerai-colbert-small-v1",
|
||||
"query": "What is machine learning?",
|
||||
"documents": [
|
||||
"Machine learning is a subset of artificial intelligence.",
|
||||
"Python is a programming language.",
|
||||
"Deep learning uses neural networks for complex tasks.",
|
||||
"The weather today is sunny.",
|
||||
],
|
||||
}
|
||||
documents = [
|
||||
"Machine learning is a subset of artificial intelligence.",
|
||||
"Python is a programming language.",
|
||||
"Deep learning uses neural networks for complex tasks.",
|
||||
"The weather today is sunny.",
|
||||
]
|
||||
|
||||
|
||||
def rerank_example():
|
||||
"""Use the /rerank endpoint to rank documents by query relevance."""
|
||||
print("=== Rerank Example ===")
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"query": "What is machine learning?",
|
||||
"documents": documents,
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
|
||||
result = response.json()
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
print("\nRanked documents (most relevant first):")
|
||||
for item in result["results"]:
|
||||
doc_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
print(f" Score {score:.4f}: {documents[doc_idx]}")
|
||||
|
||||
|
||||
def score_example():
|
||||
"""Use the /score endpoint for pairwise query-document scoring."""
|
||||
print("\n=== Score Example ===")
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"text_1": "What is machine learning?",
|
||||
"text_2": [
|
||||
"Machine learning is a subset of AI.",
|
||||
"The weather is sunny.",
|
||||
],
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
|
||||
result = response.json()
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
|
||||
def main():
|
||||
response = requests.post(url, headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
print("ColBERT Rerank Request successful!")
|
||||
result = response.json()
|
||||
print(json.dumps(result, indent=2))
|
||||
|
||||
# Show ranked results
|
||||
print("\nRanked documents (most relevant first):")
|
||||
for item in result["results"]:
|
||||
doc_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
print(f" Score {score:.4f}: {data['documents'][doc_idx]}")
|
||||
else:
|
||||
print(f"Request failed with status code: {response.status_code}")
|
||||
print(response.text)
|
||||
rerank_example()
|
||||
score_example()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Example of using ColQwen3 late interaction model for reranking.
|
||||
|
||||
ColQwen3 is a multi-modal ColBERT-style model based on Qwen3-VL.
|
||||
It produces per-token embeddings and uses MaxSim scoring for retrieval
|
||||
and reranking. Supports both text and image inputs.
|
||||
|
||||
Start the server with:
|
||||
vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 50000
|
||||
|
||||
Then run this script:
|
||||
python colqwen3_rerank_online.py
|
||||
"""
|
||||
|
||||
import requests
|
||||
|
||||
MODEL = "TomoroAI/tomoro-colqwen3-embed-4b"
|
||||
BASE_URL = "http://127.0.0.1:8000"
|
||||
|
||||
headers = {"accept": "application/json", "Content-Type": "application/json"}
|
||||
|
||||
|
||||
def rerank_text():
|
||||
"""Text-only reranking via /rerank endpoint."""
|
||||
print("=" * 60)
|
||||
print("1. Text reranking (/rerank)")
|
||||
print("=" * 60)
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"query": "What is machine learning?",
|
||||
"documents": [
|
||||
"Machine learning is a subset of artificial intelligence.",
|
||||
"Python is a programming language.",
|
||||
"Deep learning uses neural networks for complex tasks.",
|
||||
"The weather today is sunny.",
|
||||
],
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print("\n Ranked documents (most relevant first):")
|
||||
for item in result["results"]:
|
||||
doc_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
print(f" [{score:.4f}] {data['documents'][doc_idx]}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
def score_text():
|
||||
"""Text-only scoring via /score endpoint."""
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("2. Text scoring (/score)")
|
||||
print("=" * 60)
|
||||
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of France is Paris.",
|
||||
"Berlin is the capital of Germany.",
|
||||
"Python is a programming language.",
|
||||
]
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"text_1": query,
|
||||
"text_2": documents,
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print(f"\n Query: {query}\n")
|
||||
for item in result["data"]:
|
||||
idx = item["index"]
|
||||
score = item["score"]
|
||||
print(f" Doc {idx} (score={score:.4f}): {documents[idx]}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
def score_text_top_n():
|
||||
"""Text reranking with top_n filtering via /rerank endpoint."""
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("3. Text reranking with top_n=2 (/rerank)")
|
||||
print("=" * 60)
|
||||
|
||||
data = {
|
||||
"model": MODEL,
|
||||
"query": "What is the capital of France?",
|
||||
"documents": [
|
||||
"The capital of France is Paris.",
|
||||
"Berlin is the capital of Germany.",
|
||||
"Python is a programming language.",
|
||||
"The Eiffel Tower is in Paris.",
|
||||
],
|
||||
"top_n": 2,
|
||||
}
|
||||
|
||||
response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
|
||||
|
||||
if response.status_code == 200:
|
||||
result = response.json()
|
||||
print(f"\n Top {data['top_n']} results:")
|
||||
for item in result["results"]:
|
||||
doc_idx = item["index"]
|
||||
score = item["relevance_score"]
|
||||
print(f" [{score:.4f}] {data['documents'][doc_idx]}")
|
||||
else:
|
||||
print(f" Request failed: {response.status_code}")
|
||||
print(f" {response.text[:300]}")
|
||||
|
||||
|
||||
def main():
|
||||
rerank_text()
|
||||
score_text()
|
||||
score_text_top_n()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,198 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
"""
|
||||
Example online usage of Pooling API for ColQwen3 multi-vector retrieval.
|
||||
|
||||
ColQwen3 is a multi-modal late interaction model based on Qwen3-VL that
|
||||
produces per-token embeddings (320-dim, L2-normalized) for both text and
|
||||
image inputs. Similarity is computed via MaxSim scoring.
|
||||
|
||||
This example mirrors the official TomoroAI inference code
|
||||
(https://huggingface.co/TomoroAI/tomoro-colqwen3-embed-4b) but uses the
|
||||
vLLM serving API instead of local HuggingFace model loading.
|
||||
|
||||
Start the server with:
|
||||
vllm serve TomoroAI/tomoro-colqwen3-embed-4b --max-model-len 4096
|
||||
|
||||
Then run this script:
|
||||
python colqwen3_token_embed_online.py
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
# ── Helpers ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def post_http_request(payload: dict, api_url: str) -> requests.Response:
|
||||
headers = {"User-Agent": "Test Client"}
|
||||
return requests.post(api_url, headers=headers, json=payload)
|
||||
|
||||
|
||||
def load_image(url: str) -> Image.Image:
|
||||
"""Download an image from URL (handles Wikimedia 403)."""
|
||||
for hdrs in ({}, {"User-Agent": "Mozilla/5.0 (compatible; ColQwen3-demo/1.0)"}):
|
||||
resp = requests.get(url, headers=hdrs, timeout=10)
|
||||
if resp.status_code == 403:
|
||||
continue
|
||||
resp.raise_for_status()
|
||||
return Image.open(BytesIO(resp.content)).convert("RGB")
|
||||
raise RuntimeError(f"Could not fetch image from {url}")
|
||||
|
||||
|
||||
def encode_image_base64(image: Image.Image) -> str:
|
||||
"""Encode a PIL image to a base64 data URI."""
|
||||
buf = BytesIO()
|
||||
image.save(buf, format="PNG")
|
||||
return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
|
||||
|
||||
|
||||
def compute_maxsim(q_emb: np.ndarray, d_emb: np.ndarray) -> float:
|
||||
"""Compute ColBERT-style MaxSim score between query and document."""
|
||||
sim = q_emb @ d_emb.T
|
||||
return float(sim.max(axis=-1).sum())
|
||||
|
||||
|
||||
# ── Encode functions ────────────────────────────────────────
|
||||
|
||||
|
||||
def encode_queries(texts: list[str], model: str, api_url: str) -> list[np.ndarray]:
|
||||
"""Encode text queries → list of multi-vector embeddings."""
|
||||
resp = post_http_request({"model": model, "input": texts}, api_url)
|
||||
return [np.array(item["data"]) for item in resp.json()["data"]]
|
||||
|
||||
|
||||
def encode_images(image_urls: list[str], model: str, api_url: str) -> list[np.ndarray]:
|
||||
"""Encode image documents → list of multi-vector embeddings.
|
||||
|
||||
Images are sent via the chat-style `messages` field so that the
|
||||
vLLM multimodal processor handles them correctly.
|
||||
"""
|
||||
embeddings = []
|
||||
for url in image_urls:
|
||||
print(f" Loading: {url.split('/')[-1]}...")
|
||||
image = load_image(url)
|
||||
image_uri = encode_image_base64(image)
|
||||
resp = post_http_request(
|
||||
{
|
||||
"model": model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": image_uri}},
|
||||
{"type": "text", "text": "Describe the image."},
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
api_url,
|
||||
)
|
||||
result = resp.json()
|
||||
if resp.status_code != 200 or "data" not in result:
|
||||
print(f" Error ({resp.status_code}): {str(result)[:200]}")
|
||||
continue
|
||||
embeddings.append(np.array(result["data"][0]["data"]))
|
||||
return embeddings
|
||||
|
||||
|
||||
# ── Main ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="TomoroAI/tomoro-colqwen3-embed-4b",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
pooling_url = f"http://{args.host}:{args.port}/pooling"
|
||||
score_url = f"http://{args.host}:{args.port}/score"
|
||||
model = args.model
|
||||
|
||||
# Same sample data as the official TomoroAI example
|
||||
queries = [
|
||||
"Retrieve the city of Singapore",
|
||||
"Retrieve the city of Beijing",
|
||||
"Retrieve the city of London",
|
||||
]
|
||||
image_urls = [
|
||||
"https://upload.wikimedia.org/wikipedia/commons/2/27/Singapore_skyline_2022.jpg",
|
||||
"https://upload.wikimedia.org/wikipedia/commons/6/61/Beijing_skyline_at_night.JPG",
|
||||
"https://upload.wikimedia.org/wikipedia/commons/4/49/London_skyline.jpg",
|
||||
]
|
||||
|
||||
# ── 1) Text query embeddings ────────────────────────────
|
||||
print("=" * 60)
|
||||
print("1. Encode text queries (multi-vector)")
|
||||
print("=" * 60)
|
||||
query_embeddings = encode_queries(queries, model, pooling_url)
|
||||
for i, emb in enumerate(query_embeddings):
|
||||
norm = float(np.linalg.norm(emb[0]))
|
||||
print(f' Query {i}: {emb.shape} (L2 norm: {norm:.4f}) "{queries[i]}"')
|
||||
|
||||
# ── 2) Image document embeddings ────────────────────────
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("2. Encode image documents (multi-vector)")
|
||||
print("=" * 60)
|
||||
doc_embeddings = encode_images(image_urls, model, pooling_url)
|
||||
for i, emb in enumerate(doc_embeddings):
|
||||
print(f" Doc {i}: {emb.shape} {image_urls[i].split('/')[-1]}")
|
||||
|
||||
# ── 3) Cross-modal MaxSim scoring ───────────────────────
|
||||
if doc_embeddings:
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("3. Cross-modal MaxSim scores (text queries × image docs)")
|
||||
print("=" * 60)
|
||||
# Header
|
||||
print(f"{'':>35s}", end="")
|
||||
for j in range(len(doc_embeddings)):
|
||||
print(f" Doc {j:>2d}", end="")
|
||||
print()
|
||||
# Score matrix
|
||||
for i, q_emb in enumerate(query_embeddings):
|
||||
print(f" {queries[i]:<33s}", end="")
|
||||
for j, d_emb in enumerate(doc_embeddings):
|
||||
score = compute_maxsim(q_emb, d_emb)
|
||||
print(f" {score:6.2f}", end="")
|
||||
print()
|
||||
|
||||
# ── 4) Text-only /score endpoint ────────────────────────
|
||||
print()
|
||||
print("=" * 60)
|
||||
print("4. Text-only late interaction scoring (/score endpoint)")
|
||||
print("=" * 60)
|
||||
text_query = "What is the capital of France?"
|
||||
text_docs = [
|
||||
"The capital of France is Paris.",
|
||||
"Berlin is the capital of Germany.",
|
||||
"Python is a programming language.",
|
||||
]
|
||||
resp = post_http_request(
|
||||
{"model": model, "text_1": text_query, "text_2": text_docs},
|
||||
score_url,
|
||||
)
|
||||
print(f' Query: "{text_query}"\n')
|
||||
for item in resp.json()["data"]:
|
||||
idx = item["index"]
|
||||
print(f" Doc {idx} (score={item['score']:.4f}): {text_docs[idx]}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
@@ -27,7 +27,7 @@ mistral_common[image,audio] >= 1.9.1 # required for voxtral test
|
||||
num2words # required for smolvlm test
|
||||
opencv-python-headless >= 4.13.0 # required for video test
|
||||
datamodel_code_generator # required for minicpm3 test
|
||||
lm-eval[api]>=0.4.9.2 # required for model evaluation test
|
||||
lm-eval[api]>=0.4.11 # required for model evaluation test
|
||||
mteb>=1.38.11, <2 # required for mteb test
|
||||
transformers==4.57.5
|
||||
tokenizers==0.22.0
|
||||
|
||||
@@ -58,7 +58,7 @@ schemathesis==3.39.15
|
||||
# OpenAI schema test
|
||||
|
||||
# Evaluation and benchmarking
|
||||
lm-eval[api]==0.4.9.2
|
||||
lm-eval[api]==0.4.11
|
||||
jiwer==4.0.0
|
||||
|
||||
# Required for multiprocessed tests that use spawn method, Datasets and Evaluate Test
|
||||
@@ -67,8 +67,6 @@ multiprocess==0.70.16
|
||||
# Required for v1/metrics/test_engine_logger_apis.py
|
||||
ray[cgraph,default]>=2.48.0
|
||||
|
||||
# Plugins test
|
||||
terratorch @ git+https://github.com/IBM/terratorch.git@07184fcf91a1324f831ff521dd238d97fe350e3e
|
||||
torchgeo==0.7.0
|
||||
# via terratorch
|
||||
# MTEB Benchmark Test
|
||||
@@ -96,3 +94,11 @@ albumentations==1.4.6
|
||||
transformers==4.57.3
|
||||
# Pin HF Hub version
|
||||
huggingface-hub==0.36.2
|
||||
# Pin Mistral Common
|
||||
mistral-common[image,audio]==1.9.1
|
||||
# Required for Prithvi tests
|
||||
terratorch==1.2.2
|
||||
# Required for Prithvi tests
|
||||
segmentation-models-pytorch==0.5.0
|
||||
# Required for Prithvi tests
|
||||
imagehash==4.3.2
|
||||
|
||||
@@ -35,7 +35,7 @@ num2words # required for smolvlm test
|
||||
open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
|
||||
opencv-python-headless >= 4.13.0 # required for video test
|
||||
datamodel_code_generator # required for minicpm3 test
|
||||
lm-eval[api]>=0.4.9.2 # required for model evaluation test
|
||||
lm-eval[api]>=0.4.11 # required for model evaluation test
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
transformers==4.57.5
|
||||
tokenizers==0.22.0
|
||||
@@ -56,7 +56,14 @@ runai-model-streamer[s3,gcs]==0.15.3
|
||||
fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
decord==0.6.0
|
||||
terratorch @ git+https://github.com/IBM/terratorch.git@1.1.rc3 # required for PrithviMAE test
|
||||
terratorch >= 1.2.2 # Required for Prithvi tests
|
||||
imagehash # Required for Prithvi tests
|
||||
segmentation-models-pytorch > 0.4.0 # Required for Prithvi tests
|
||||
|
||||
gpt-oss >= 0.0.7; python_version > '3.11'
|
||||
|
||||
perceptron # required for isaac test
|
||||
|
||||
# Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library.
|
||||
# Older versions are in conflict with teerratorch requirements.
|
||||
datasets>=3.3.0,<=3.6.0
|
||||
+108
-127
@@ -1,11 +1,11 @@
|
||||
# This file was autogenerated by uv via the following command:
|
||||
# uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu129 --python-platform x86_64-manylinux_2_28 --python-version 3.12
|
||||
# uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu128 --python-platform x86_64-manylinux_2_28 --python-version 3.12
|
||||
absl-py==2.1.0
|
||||
# via rouge-score
|
||||
accelerate==1.0.1
|
||||
# via
|
||||
# lm-eval
|
||||
# peft
|
||||
# rouge-score
|
||||
# tensorboard
|
||||
accelerate==1.0.1
|
||||
# via peft
|
||||
aenum==3.1.16
|
||||
# via lightly
|
||||
affine==2.4.0
|
||||
@@ -31,9 +31,7 @@ albumentations==1.4.6
|
||||
# -r requirements/test.in
|
||||
# terratorch
|
||||
alembic==1.16.4
|
||||
# via
|
||||
# mlflow
|
||||
# optuna
|
||||
# via optuna
|
||||
annotated-doc==0.0.4
|
||||
# via fastapi
|
||||
annotated-types==0.7.0
|
||||
@@ -74,8 +72,6 @@ bitsandbytes==0.46.1
|
||||
# lightning
|
||||
black==24.10.0
|
||||
# via datamodel-code-generator
|
||||
blinker==1.9.0
|
||||
# via flask
|
||||
blobfile==3.0.0
|
||||
# via -r requirements/test.in
|
||||
bm25s==0.2.13
|
||||
@@ -93,9 +89,7 @@ bounded-pool-executor==0.0.3
|
||||
buildkite-test-collector==0.1.9
|
||||
# via -r requirements/test.in
|
||||
cachetools==5.5.2
|
||||
# via
|
||||
# google-auth
|
||||
# mlflow-skinny
|
||||
# via google-auth
|
||||
certifi==2024.8.30
|
||||
# via
|
||||
# fiona
|
||||
@@ -106,6 +100,7 @@ certifi==2024.8.30
|
||||
# pyproj
|
||||
# rasterio
|
||||
# requests
|
||||
# sentry-sdk
|
||||
cffi==1.17.1
|
||||
# via soundfile
|
||||
chardet==5.2.0
|
||||
@@ -120,15 +115,14 @@ click==8.1.7
|
||||
# click-plugins
|
||||
# cligj
|
||||
# fiona
|
||||
# flask
|
||||
# jiwer
|
||||
# mlflow-skinny
|
||||
# nltk
|
||||
# rasterio
|
||||
# ray
|
||||
# schemathesis
|
||||
# typer
|
||||
# uvicorn
|
||||
# wandb
|
||||
click-plugins==1.1.1.2
|
||||
# via
|
||||
# fiona
|
||||
@@ -137,14 +131,11 @@ cligj==0.7.2
|
||||
# via
|
||||
# fiona
|
||||
# rasterio
|
||||
cloudpickle==3.1.1
|
||||
# via mlflow-skinny
|
||||
colorama==0.4.6
|
||||
# via
|
||||
# perceptron
|
||||
# sacrebleu
|
||||
# schemathesis
|
||||
# tqdm-multiprocess
|
||||
colorful==0.5.6
|
||||
# via ray
|
||||
colorlog==6.10.1
|
||||
@@ -163,16 +154,15 @@ cupy-cuda12x==13.6.0
|
||||
# via ray
|
||||
cycler==0.12.1
|
||||
# via matplotlib
|
||||
databricks-sdk==0.59.0
|
||||
# via mlflow-skinny
|
||||
datamodel-code-generator==0.26.3
|
||||
# via -r requirements/test.in
|
||||
dataproperty==1.0.1
|
||||
# via
|
||||
# pytablewriter
|
||||
# tabledata
|
||||
datasets==3.0.2
|
||||
datasets==3.3.0
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# evaluate
|
||||
# lm-eval
|
||||
# mteb
|
||||
@@ -180,6 +170,8 @@ decorator==5.1.1
|
||||
# via librosa
|
||||
decord==0.6.0
|
||||
# via -r requirements/test.in
|
||||
diffusers==0.36.0
|
||||
# via terratorch
|
||||
dill==0.3.8
|
||||
# via
|
||||
# datasets
|
||||
@@ -191,15 +183,11 @@ distlib==0.3.9
|
||||
dnspython==2.7.0
|
||||
# via email-validator
|
||||
docker==7.1.0
|
||||
# via
|
||||
# gpt-oss
|
||||
# mlflow
|
||||
# via gpt-oss
|
||||
docopt==0.6.2
|
||||
# via num2words
|
||||
docstring-parser==0.17.0
|
||||
# via jsonargparse
|
||||
efficientnet-pytorch==0.7.1
|
||||
# via segmentation-models-pytorch
|
||||
einops==0.8.1
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
@@ -217,9 +205,7 @@ encodec==0.1.1
|
||||
evaluate==0.4.3
|
||||
# via lm-eval
|
||||
fastapi==0.128.0
|
||||
# via
|
||||
# gpt-oss
|
||||
# mlflow-skinny
|
||||
# via gpt-oss
|
||||
fastparquet==2024.11.0
|
||||
# via genai-perf
|
||||
fastrlock==0.8.2
|
||||
@@ -230,6 +216,7 @@ filelock==3.16.1
|
||||
# via
|
||||
# blobfile
|
||||
# datasets
|
||||
# diffusers
|
||||
# huggingface-hub
|
||||
# ray
|
||||
# torch
|
||||
@@ -237,8 +224,6 @@ filelock==3.16.1
|
||||
# virtualenv
|
||||
fiona==1.10.1
|
||||
# via torchgeo
|
||||
flask==3.1.1
|
||||
# via mlflow
|
||||
fonttools==4.55.0
|
||||
# via matplotlib
|
||||
fqdn==1.5.1
|
||||
@@ -249,7 +234,7 @@ frozenlist==1.5.0
|
||||
# via
|
||||
# aiohttp
|
||||
# aiosignal
|
||||
fsspec==2024.9.0
|
||||
fsspec==2024.12.0
|
||||
# via
|
||||
# datasets
|
||||
# evaluate
|
||||
@@ -257,6 +242,7 @@ fsspec==2024.9.0
|
||||
# huggingface-hub
|
||||
# lightning
|
||||
# pytorch-lightning
|
||||
# tacoreader
|
||||
# torch
|
||||
ftfy==6.3.1
|
||||
# via open-clip-torch
|
||||
@@ -269,7 +255,7 @@ geopandas==1.0.1
|
||||
gitdb==4.0.12
|
||||
# via gitpython
|
||||
gitpython==3.1.44
|
||||
# via mlflow-skinny
|
||||
# via wandb
|
||||
google-api-core==2.24.2
|
||||
# via
|
||||
# google-cloud-core
|
||||
@@ -277,7 +263,6 @@ google-api-core==2.24.2
|
||||
# opencensus
|
||||
google-auth==2.40.2
|
||||
# via
|
||||
# databricks-sdk
|
||||
# google-api-core
|
||||
# google-cloud-core
|
||||
# google-cloud-storage
|
||||
@@ -296,25 +281,17 @@ googleapis-common-protos==1.70.0
|
||||
# via google-api-core
|
||||
gpt-oss==0.0.8
|
||||
# via -r requirements/test.in
|
||||
graphene==3.4.3
|
||||
# via mlflow
|
||||
graphql-core==3.2.6
|
||||
# via
|
||||
# graphene
|
||||
# graphql-relay
|
||||
# hypothesis-graphql
|
||||
graphql-relay==3.2.0
|
||||
# via graphene
|
||||
# via hypothesis-graphql
|
||||
greenlet==3.2.3
|
||||
# via sqlalchemy
|
||||
grpcio==1.78.0
|
||||
# via
|
||||
# grpcio-tools
|
||||
# ray
|
||||
# tensorboard
|
||||
grpcio-tools==1.78.0
|
||||
# via -r requirements/test.in
|
||||
gunicorn==23.0.0
|
||||
# via mlflow
|
||||
h11==0.14.0
|
||||
# via
|
||||
# httpcore
|
||||
@@ -338,12 +315,14 @@ httpcore==1.0.6
|
||||
httpx==0.27.2
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# diffusers
|
||||
# perceptron
|
||||
# schemathesis
|
||||
huggingface-hub==0.36.2
|
||||
# via
|
||||
# accelerate
|
||||
# datasets
|
||||
# diffusers
|
||||
# evaluate
|
||||
# open-clip-torch
|
||||
# peft
|
||||
@@ -379,11 +358,13 @@ idna==3.10
|
||||
# jsonschema
|
||||
# requests
|
||||
# yarl
|
||||
imagehash==4.3.2
|
||||
# via -r requirements/test.in
|
||||
imageio==2.37.0
|
||||
# via scikit-image
|
||||
importlib-metadata==8.7.0
|
||||
# via
|
||||
# mlflow-skinny
|
||||
# diffusers
|
||||
# opentelemetry-api
|
||||
importlib-resources==6.5.2
|
||||
# via typeshed-client
|
||||
@@ -395,14 +376,11 @@ isoduration==20.11.0
|
||||
# via jsonschema
|
||||
isort==5.13.2
|
||||
# via datamodel-code-generator
|
||||
itsdangerous==2.2.0
|
||||
# via flask
|
||||
jinja2==3.1.6
|
||||
# via
|
||||
# datamodel-code-generator
|
||||
# flask
|
||||
# genai-perf
|
||||
# mlflow
|
||||
# lm-eval
|
||||
# torch
|
||||
jiwer==3.0.5
|
||||
# via -r requirements/test.in
|
||||
@@ -415,12 +393,14 @@ joblib==1.4.2
|
||||
# librosa
|
||||
# nltk
|
||||
# scikit-learn
|
||||
jsonargparse==4.35.0
|
||||
jsonargparse==4.46.0
|
||||
# via
|
||||
# lightning
|
||||
# terratorch
|
||||
jsonlines==4.0.0
|
||||
# via lm-eval
|
||||
jsonnet==0.21.0
|
||||
# via jsonargparse
|
||||
jsonpointer==3.0.0
|
||||
# via jsonschema
|
||||
jsonschema==4.23.0
|
||||
@@ -449,13 +429,13 @@ libnacl==2.1.0
|
||||
# via tensorizer
|
||||
librosa==0.10.2.post1
|
||||
# via -r requirements/test.in
|
||||
lightly==1.5.20
|
||||
lightly==1.5.22
|
||||
# via
|
||||
# terratorch
|
||||
# torchgeo
|
||||
lightly-utils==0.0.2
|
||||
# via lightly
|
||||
lightning==2.5.1.post0
|
||||
lightning==2.6.1
|
||||
# via
|
||||
# terratorch
|
||||
# torchgeo
|
||||
@@ -466,7 +446,7 @@ lightning-utilities==0.14.3
|
||||
# torchmetrics
|
||||
llvmlite==0.44.0
|
||||
# via numba
|
||||
lm-eval==0.4.9.2
|
||||
lm-eval==0.4.11
|
||||
# via -r requirements/test.in
|
||||
lxml==5.3.0
|
||||
# via
|
||||
@@ -476,12 +456,11 @@ lxml==5.3.0
|
||||
mako==1.3.10
|
||||
# via alembic
|
||||
markdown==3.8.2
|
||||
# via mlflow
|
||||
# via tensorboard
|
||||
markdown-it-py==3.0.0
|
||||
# via rich
|
||||
markupsafe==3.0.1
|
||||
# via
|
||||
# flask
|
||||
# jinja2
|
||||
# mako
|
||||
# werkzeug
|
||||
@@ -489,7 +468,6 @@ matplotlib==3.9.2
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# lightning
|
||||
# mlflow
|
||||
# pycocotools
|
||||
# torchgeo
|
||||
mbstrdecoder==1.1.3
|
||||
@@ -501,10 +479,6 @@ mdurl==0.1.2
|
||||
# via markdown-it-py
|
||||
mistral-common==1.9.1
|
||||
# via -r requirements/test.in
|
||||
mlflow==2.22.0
|
||||
# via terratorch
|
||||
mlflow-skinny==2.22.0
|
||||
# via mlflow
|
||||
more-itertools==10.5.0
|
||||
# via lm-eval
|
||||
mpmath==1.3.0
|
||||
@@ -523,8 +497,6 @@ multiprocess==0.70.16
|
||||
# via
|
||||
# datasets
|
||||
# evaluate
|
||||
munch==4.0.0
|
||||
# via pretrainedmodels
|
||||
mypy-extensions==1.0.0
|
||||
# via black
|
||||
networkx==3.2.1
|
||||
@@ -539,8 +511,6 @@ numba==0.61.2
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# librosa
|
||||
numexpr==2.10.1
|
||||
# via lm-eval
|
||||
numpy==2.2.6
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
@@ -553,6 +523,7 @@ numpy==2.2.6
|
||||
# cupy-cuda12x
|
||||
# datasets
|
||||
# decord
|
||||
# diffusers
|
||||
# einx
|
||||
# encodec
|
||||
# evaluate
|
||||
@@ -560,16 +531,16 @@ numpy==2.2.6
|
||||
# genai-perf
|
||||
# geopandas
|
||||
# h5py
|
||||
# imagehash
|
||||
# imageio
|
||||
# librosa
|
||||
# lightly
|
||||
# lightly-utils
|
||||
# lm-eval
|
||||
# matplotlib
|
||||
# mistral-common
|
||||
# mlflow
|
||||
# mteb
|
||||
# numba
|
||||
# numexpr
|
||||
# opencv-python-headless
|
||||
# optuna
|
||||
# pandas
|
||||
@@ -578,6 +549,7 @@ numpy==2.2.6
|
||||
# perceptron
|
||||
# pycocotools
|
||||
# pyogrio
|
||||
# pywavelets
|
||||
# rasterio
|
||||
# rioxarray
|
||||
# rouge-score
|
||||
@@ -590,8 +562,10 @@ numpy==2.2.6
|
||||
# shapely
|
||||
# soxr
|
||||
# statsmodels
|
||||
# tensorboard
|
||||
# tensorboardx
|
||||
# tensorizer
|
||||
# terratorch
|
||||
# tifffile
|
||||
# torchgeo
|
||||
# torchmetrics
|
||||
@@ -600,28 +574,28 @@ numpy==2.2.6
|
||||
# tritonclient
|
||||
# vocos
|
||||
# xarray
|
||||
nvidia-cublas-cu12==12.9.1.4
|
||||
nvidia-cublas-cu12==12.8.4.1
|
||||
# via
|
||||
# nvidia-cudnn-cu12
|
||||
# nvidia-cusolver-cu12
|
||||
# torch
|
||||
nvidia-cuda-cupti-cu12==12.9.79
|
||||
nvidia-cuda-cupti-cu12==12.8.90
|
||||
# via torch
|
||||
nvidia-cuda-nvrtc-cu12==12.9.86
|
||||
nvidia-cuda-nvrtc-cu12==12.8.93
|
||||
# via torch
|
||||
nvidia-cuda-runtime-cu12==12.9.79
|
||||
nvidia-cuda-runtime-cu12==12.8.90
|
||||
# via torch
|
||||
nvidia-cudnn-cu12==9.10.2.21
|
||||
# via torch
|
||||
nvidia-cufft-cu12==11.4.1.4
|
||||
nvidia-cufft-cu12==11.3.3.83
|
||||
# via torch
|
||||
nvidia-cufile-cu12==1.14.1.1
|
||||
nvidia-cufile-cu12==1.13.1.3
|
||||
# via torch
|
||||
nvidia-curand-cu12==10.3.10.19
|
||||
nvidia-curand-cu12==10.3.9.90
|
||||
# via torch
|
||||
nvidia-cusolver-cu12==11.7.5.82
|
||||
nvidia-cusolver-cu12==11.7.3.90
|
||||
# via torch
|
||||
nvidia-cusparse-cu12==12.5.10.65
|
||||
nvidia-cusparse-cu12==12.5.8.93
|
||||
# via
|
||||
# nvidia-cusolver-cu12
|
||||
# torch
|
||||
@@ -629,7 +603,7 @@ nvidia-cusparselt-cu12==0.7.1
|
||||
# via torch
|
||||
nvidia-nccl-cu12==2.27.5
|
||||
# via torch
|
||||
nvidia-nvjitlink-cu12==12.9.86
|
||||
nvidia-nvjitlink-cu12==12.8.93
|
||||
# via
|
||||
# nvidia-cufft-cu12
|
||||
# nvidia-cusolver-cu12
|
||||
@@ -637,7 +611,7 @@ nvidia-nvjitlink-cu12==12.9.86
|
||||
# torch
|
||||
nvidia-nvshmem-cu12==3.4.5
|
||||
# via torch
|
||||
nvidia-nvtx-cu12==12.9.79
|
||||
nvidia-nvtx-cu12==12.8.90
|
||||
# via torch
|
||||
omegaconf==2.3.0
|
||||
# via
|
||||
@@ -659,7 +633,6 @@ opencv-python-headless==4.13.0.90
|
||||
# mistral-common
|
||||
opentelemetry-api==1.35.0
|
||||
# via
|
||||
# mlflow-skinny
|
||||
# opentelemetry-exporter-prometheus
|
||||
# opentelemetry-sdk
|
||||
# opentelemetry-semantic-conventions
|
||||
@@ -669,7 +642,6 @@ opentelemetry-proto==1.36.0
|
||||
# via ray
|
||||
opentelemetry-sdk==1.35.0
|
||||
# via
|
||||
# mlflow-skinny
|
||||
# opentelemetry-exporter-prometheus
|
||||
# ray
|
||||
opentelemetry-semantic-conventions==0.56b0
|
||||
@@ -687,7 +659,6 @@ packaging==24.2
|
||||
# evaluate
|
||||
# fastparquet
|
||||
# geopandas
|
||||
# gunicorn
|
||||
# huggingface-hub
|
||||
# hydra-core
|
||||
# kornia
|
||||
@@ -695,7 +666,6 @@ packaging==24.2
|
||||
# lightning
|
||||
# lightning-utilities
|
||||
# matplotlib
|
||||
# mlflow-skinny
|
||||
# optuna
|
||||
# peft
|
||||
# plotly
|
||||
@@ -708,10 +678,12 @@ packaging==24.2
|
||||
# rioxarray
|
||||
# scikit-image
|
||||
# statsmodels
|
||||
# tensorboard
|
||||
# tensorboardx
|
||||
# torchmetrics
|
||||
# transformers
|
||||
# typepy
|
||||
# wandb
|
||||
# xarray
|
||||
pandas==2.2.3
|
||||
# via
|
||||
@@ -720,8 +692,8 @@ pandas==2.2.3
|
||||
# fastparquet
|
||||
# genai-perf
|
||||
# geopandas
|
||||
# mlflow
|
||||
# statsmodels
|
||||
# tacoreader
|
||||
# torchgeo
|
||||
# xarray
|
||||
pathspec==0.12.1
|
||||
@@ -731,16 +703,16 @@ pathvalidate==3.2.1
|
||||
patsy==1.0.1
|
||||
# via statsmodels
|
||||
peft==0.16.0
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# lm-eval
|
||||
# via -r requirements/test.in
|
||||
perceptron==0.1.4
|
||||
# via -r requirements/test.in
|
||||
perf-analyzer==0.1.0
|
||||
# via genai-perf
|
||||
pillow==10.4.0
|
||||
# via
|
||||
# diffusers
|
||||
# genai-perf
|
||||
# imagehash
|
||||
# imageio
|
||||
# lightly-utils
|
||||
# matplotlib
|
||||
@@ -748,6 +720,7 @@ pillow==10.4.0
|
||||
# perceptron
|
||||
# scikit-image
|
||||
# segmentation-models-pytorch
|
||||
# tensorboard
|
||||
# torchgeo
|
||||
# torchvision
|
||||
platformdirs==4.3.6
|
||||
@@ -755,6 +728,7 @@ platformdirs==4.3.6
|
||||
# black
|
||||
# pooch
|
||||
# virtualenv
|
||||
# wandb
|
||||
plotly==5.24.1
|
||||
# via genai-perf
|
||||
pluggy==1.5.0
|
||||
@@ -769,8 +743,6 @@ portalocker==2.10.1
|
||||
# via sacrebleu
|
||||
pqdm==0.2.0
|
||||
# via -r requirements/test.in
|
||||
pretrainedmodels==0.7.4
|
||||
# via segmentation-models-pytorch
|
||||
prometheus-client==0.22.0
|
||||
# via
|
||||
# opentelemetry-exporter-prometheus
|
||||
@@ -786,12 +758,13 @@ protobuf==6.33.2
|
||||
# google-api-core
|
||||
# googleapis-common-protos
|
||||
# grpcio-tools
|
||||
# mlflow-skinny
|
||||
# opentelemetry-proto
|
||||
# proto-plus
|
||||
# ray
|
||||
# tensorboard
|
||||
# tensorboardx
|
||||
# tensorizer
|
||||
# wandb
|
||||
psutil==6.1.0
|
||||
# via
|
||||
# accelerate
|
||||
@@ -801,19 +774,18 @@ py==1.11.0
|
||||
# via pytest-forked
|
||||
py-spy==0.4.0
|
||||
# via ray
|
||||
pyarrow==18.0.0
|
||||
pyarrow==23.0.0
|
||||
# via
|
||||
# datasets
|
||||
# genai-perf
|
||||
# mlflow
|
||||
# tacoreader
|
||||
# terratorch
|
||||
pyasn1==0.6.1
|
||||
# via
|
||||
# pyasn1-modules
|
||||
# rsa
|
||||
pyasn1-modules==0.4.2
|
||||
# via google-auth
|
||||
pybind11==2.13.6
|
||||
# via lm-eval
|
||||
pycocotools==2.0.8
|
||||
# via terratorch
|
||||
pycountry==24.6.1
|
||||
@@ -831,11 +803,11 @@ pydantic==2.12.0
|
||||
# gpt-oss
|
||||
# lightly
|
||||
# mistral-common
|
||||
# mlflow-skinny
|
||||
# mteb
|
||||
# openai-harmony
|
||||
# pydantic-extra-types
|
||||
# ray
|
||||
# wandb
|
||||
pydantic-core==2.41.1
|
||||
# via pydantic
|
||||
pydantic-extra-types==2.10.5
|
||||
@@ -873,7 +845,6 @@ pytest==8.3.5
|
||||
# pytest-subtests
|
||||
# pytest-timeout
|
||||
# schemathesis
|
||||
# terratorch
|
||||
pytest-asyncio==0.24.0
|
||||
# via -r requirements/test.in
|
||||
pytest-cov==6.3.0
|
||||
@@ -896,7 +867,6 @@ python-dateutil==2.9.0.post0
|
||||
# via
|
||||
# arrow
|
||||
# botocore
|
||||
# graphene
|
||||
# lightly
|
||||
# matplotlib
|
||||
# pandas
|
||||
@@ -913,6 +883,8 @@ pytz==2024.2
|
||||
# via
|
||||
# pandas
|
||||
# typepy
|
||||
pywavelets==1.9.0
|
||||
# via imagehash
|
||||
pyyaml==6.0.2
|
||||
# via
|
||||
# accelerate
|
||||
@@ -923,7 +895,6 @@ pyyaml==6.0.2
|
||||
# huggingface-hub
|
||||
# jsonargparse
|
||||
# lightning
|
||||
# mlflow-skinny
|
||||
# omegaconf
|
||||
# optuna
|
||||
# peft
|
||||
@@ -934,6 +905,7 @@ pyyaml==6.0.2
|
||||
# timm
|
||||
# transformers
|
||||
# vocos
|
||||
# wandb
|
||||
rapidfuzz==3.12.1
|
||||
# via jiwer
|
||||
rasterio==1.4.3
|
||||
@@ -951,6 +923,7 @@ referencing==0.35.1
|
||||
# jsonschema-specifications
|
||||
regex==2024.9.11
|
||||
# via
|
||||
# diffusers
|
||||
# nltk
|
||||
# open-clip-torch
|
||||
# sacrebleu
|
||||
@@ -959,8 +932,8 @@ regex==2024.9.11
|
||||
requests==2.32.3
|
||||
# via
|
||||
# buildkite-test-collector
|
||||
# databricks-sdk
|
||||
# datasets
|
||||
# diffusers
|
||||
# docker
|
||||
# evaluate
|
||||
# google-api-core
|
||||
@@ -970,15 +943,16 @@ requests==2.32.3
|
||||
# lightly
|
||||
# lm-eval
|
||||
# mistral-common
|
||||
# mlflow-skinny
|
||||
# mteb
|
||||
# pooch
|
||||
# ray
|
||||
# responses
|
||||
# schemathesis
|
||||
# starlette-testclient
|
||||
# tacoreader
|
||||
# tiktoken
|
||||
# transformers
|
||||
# wandb
|
||||
responses==0.25.3
|
||||
# via genai-perf
|
||||
rfc3339-validator==0.1.4
|
||||
@@ -991,6 +965,7 @@ rich==13.9.4
|
||||
# lightning
|
||||
# mteb
|
||||
# perceptron
|
||||
# terratorch
|
||||
# typer
|
||||
rioxarray==0.19.0
|
||||
# via terratorch
|
||||
@@ -1017,47 +992,55 @@ sacrebleu==2.4.3
|
||||
safetensors==0.4.5
|
||||
# via
|
||||
# accelerate
|
||||
# diffusers
|
||||
# open-clip-torch
|
||||
# peft
|
||||
# segmentation-models-pytorch
|
||||
# timm
|
||||
# transformers
|
||||
schemathesis==3.39.15
|
||||
# via -r requirements/test.in
|
||||
scikit-image==0.25.2
|
||||
# via albumentations
|
||||
# via
|
||||
# albumentations
|
||||
# terratorch
|
||||
scikit-learn==1.5.2
|
||||
# via
|
||||
# albumentations
|
||||
# librosa
|
||||
# lm-eval
|
||||
# mlflow
|
||||
# mteb
|
||||
# sentence-transformers
|
||||
# terratorch
|
||||
scipy==1.13.1
|
||||
# via
|
||||
# albumentations
|
||||
# bm25s
|
||||
# imagehash
|
||||
# librosa
|
||||
# mlflow
|
||||
# mteb
|
||||
# scikit-image
|
||||
# scikit-learn
|
||||
# sentence-transformers
|
||||
# statsmodels
|
||||
# vocos
|
||||
segmentation-models-pytorch==0.4.0
|
||||
segmentation-models-pytorch==0.5.0
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# terratorch
|
||||
# torchgeo
|
||||
sentence-transformers==5.2.0
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# mteb
|
||||
sentry-sdk==2.52.0
|
||||
# via wandb
|
||||
setuptools==77.0.3
|
||||
# via
|
||||
# grpcio-tools
|
||||
# lightning-utilities
|
||||
# pytablewriter
|
||||
# tensorboard
|
||||
# torch
|
||||
shapely==2.1.1
|
||||
# via
|
||||
@@ -1075,7 +1058,6 @@ six==1.16.0
|
||||
# python-dateutil
|
||||
# rfc3339-validator
|
||||
# rouge-score
|
||||
# segmentation-models-pytorch
|
||||
smart-open==7.1.0
|
||||
# via ray
|
||||
smmap==5.0.2
|
||||
@@ -1099,12 +1081,9 @@ soxr==0.5.0.post1
|
||||
sqlalchemy==2.0.41
|
||||
# via
|
||||
# alembic
|
||||
# mlflow
|
||||
# optuna
|
||||
sqlitedict==2.1.0
|
||||
# via lm-eval
|
||||
sqlparse==0.5.3
|
||||
# via mlflow-skinny
|
||||
starlette==0.50.0
|
||||
# via
|
||||
# fastapi
|
||||
@@ -1124,6 +1103,8 @@ tabledata==1.3.3
|
||||
# via pytablewriter
|
||||
tabulate==0.9.0
|
||||
# via sacrebleu
|
||||
tacoreader==0.5.6
|
||||
# via terratorch
|
||||
tblib==3.1.0
|
||||
# via -r requirements/test.in
|
||||
tcolorpy==0.1.6
|
||||
@@ -1133,13 +1114,19 @@ tenacity==9.1.2
|
||||
# gpt-oss
|
||||
# lm-eval
|
||||
# plotly
|
||||
tensorboard==2.20.0
|
||||
# via terratorch
|
||||
tensorboard-data-server==0.7.2
|
||||
# via tensorboard
|
||||
tensorboardx==2.6.4
|
||||
# via lightning
|
||||
tensorizer==2.10.1
|
||||
# via -r requirements/test.in
|
||||
termcolor==3.1.0
|
||||
# via gpt-oss
|
||||
terratorch @ git+https://github.com/IBM/terratorch.git@07184fcf91a1324f831ff521dd238d97fe350e3e
|
||||
# via
|
||||
# gpt-oss
|
||||
# terratorch
|
||||
terratorch==1.2.2
|
||||
# via -r requirements/test.in
|
||||
threadpoolctl==3.5.0
|
||||
# via scikit-learn
|
||||
@@ -1167,21 +1154,18 @@ tomli==2.2.1
|
||||
# via schemathesis
|
||||
tomli-w==1.2.0
|
||||
# via schemathesis
|
||||
torch==2.10.0+cu129
|
||||
torch==2.10.0+cu128
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# accelerate
|
||||
# bitsandbytes
|
||||
# efficientnet-pytorch
|
||||
# encodec
|
||||
# kornia
|
||||
# lightly
|
||||
# lightning
|
||||
# lm-eval
|
||||
# mteb
|
||||
# open-clip-torch
|
||||
# peft
|
||||
# pretrainedmodels
|
||||
# pytorch-lightning
|
||||
# runai-model-streamer
|
||||
# segmentation-models-pytorch
|
||||
@@ -1195,7 +1179,7 @@ torch==2.10.0+cu129
|
||||
# torchvision
|
||||
# vector-quantize-pytorch
|
||||
# vocos
|
||||
torchaudio==2.10.0+cu129
|
||||
torchaudio==2.10.0+cu128
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# encodec
|
||||
@@ -1208,17 +1192,16 @@ torchmetrics==1.7.4
|
||||
# pytorch-lightning
|
||||
# terratorch
|
||||
# torchgeo
|
||||
torchvision==0.25.0+cu129
|
||||
torchvision==0.25.0+cu128
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# lightly
|
||||
# open-clip-torch
|
||||
# pretrainedmodels
|
||||
# segmentation-models-pytorch
|
||||
# terratorch
|
||||
# timm
|
||||
# torchgeo
|
||||
tqdm==4.66.6
|
||||
tqdm==4.67.3
|
||||
# via
|
||||
# datasets
|
||||
# evaluate
|
||||
@@ -1232,19 +1215,16 @@ tqdm==4.66.6
|
||||
# optuna
|
||||
# peft
|
||||
# pqdm
|
||||
# pretrainedmodels
|
||||
# pytorch-lightning
|
||||
# segmentation-models-pytorch
|
||||
# sentence-transformers
|
||||
# tqdm-multiprocess
|
||||
# tacoreader
|
||||
# terratorch
|
||||
# transformers
|
||||
tqdm-multiprocess==0.0.11
|
||||
# via lm-eval
|
||||
transformers==4.57.5
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# genai-perf
|
||||
# lm-eval
|
||||
# peft
|
||||
# sentence-transformers
|
||||
# transformers-stream-generator
|
||||
@@ -1274,14 +1254,13 @@ typing-extensions==4.15.0
|
||||
# alembic
|
||||
# chz
|
||||
# fastapi
|
||||
# graphene
|
||||
# grpcio
|
||||
# huggingface-hub
|
||||
# librosa
|
||||
# lightning
|
||||
# lightning-utilities
|
||||
# lm-eval
|
||||
# mistral-common
|
||||
# mlflow-skinny
|
||||
# mteb
|
||||
# opentelemetry-api
|
||||
# opentelemetry-sdk
|
||||
@@ -1299,6 +1278,7 @@ typing-extensions==4.15.0
|
||||
# typer
|
||||
# typeshed-client
|
||||
# typing-inspection
|
||||
# wandb
|
||||
typing-inspection==0.4.2
|
||||
# via pydantic
|
||||
tzdata==2024.2
|
||||
@@ -1313,25 +1293,26 @@ urllib3==2.2.3
|
||||
# lightly
|
||||
# requests
|
||||
# responses
|
||||
# sentry-sdk
|
||||
# tritonclient
|
||||
uvicorn==0.35.0
|
||||
# via
|
||||
# gpt-oss
|
||||
# mlflow-skinny
|
||||
# via gpt-oss
|
||||
vector-quantize-pytorch==1.21.2
|
||||
# via -r requirements/test.in
|
||||
virtualenv==20.31.2
|
||||
# via ray
|
||||
vocos==0.1.0
|
||||
# via -r requirements/test.in
|
||||
wandb==0.24.2
|
||||
# via terratorch
|
||||
wcwidth==0.2.13
|
||||
# via ftfy
|
||||
webcolors==24.11.1
|
||||
# via jsonschema
|
||||
werkzeug==3.1.3
|
||||
# via
|
||||
# flask
|
||||
# schemathesis
|
||||
# tensorboard
|
||||
word2number==1.1
|
||||
# via lm-eval
|
||||
wrapt==1.17.2
|
||||
|
||||
@@ -0,0 +1,430 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Autotune registered Helion kernels for optimal configurations.
|
||||
|
||||
Usage:
|
||||
# Autotune all registered kernels
|
||||
python scripts/autotune_helion_kernels.py
|
||||
|
||||
# Autotune specific kernel
|
||||
python scripts/autotune_helion_kernels.py --kernels silu_mul_fp8
|
||||
|
||||
# Autotune multiple kernels
|
||||
python scripts/autotune_helion_kernels.py --kernels silu_mul_fp8 rms_norm_fp8
|
||||
|
||||
# Force re-autotuning
|
||||
python scripts/autotune_helion_kernels.py --force
|
||||
|
||||
# List available kernels
|
||||
python scripts/autotune_helion_kernels.py --list
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
try:
|
||||
import helion
|
||||
|
||||
from vllm.kernels.helion import (
|
||||
ConfigManager,
|
||||
get_kernel_by_name,
|
||||
get_registered_kernels,
|
||||
)
|
||||
from vllm.kernels.helion.utils import get_canonical_gpu_name
|
||||
from vllm.logger import init_logger
|
||||
from vllm.utils.import_utils import has_helion
|
||||
except ImportError as e:
|
||||
print(f"Error importing vLLM: {e}")
|
||||
print("Please ensure vLLM is installed and in your Python path")
|
||||
sys.exit(1)
|
||||
|
||||
logger = init_logger("vllm.scripts.autotune_helion_kernels")
|
||||
|
||||
|
||||
@dataclass
|
||||
class AutotuneResult:
|
||||
status: str # "success" | "partial" | "error" | "skipped"
|
||||
successful: int
|
||||
failed: int
|
||||
configs: dict[str, "helion.Config"]
|
||||
message: str = ""
|
||||
|
||||
|
||||
def list_kernels() -> None:
|
||||
kernels = get_registered_kernels()
|
||||
|
||||
if not kernels:
|
||||
print("No Helion kernels found in registry.")
|
||||
return
|
||||
|
||||
print("Available Helion kernels:")
|
||||
print("=" * 50)
|
||||
|
||||
for name in sorted(kernels.keys()):
|
||||
print(f" {name}")
|
||||
|
||||
print(f"\nTotal: {len(kernels)} kernels")
|
||||
|
||||
|
||||
def check_requirements() -> bool:
|
||||
if not torch.cuda.is_available():
|
||||
logger.error("CUDA is not available. Helion autotuning requires GPU.")
|
||||
return False
|
||||
|
||||
if not has_helion():
|
||||
logger.error("Helion is not installed. Please install Helion package.")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def autotune_kernel(
|
||||
kernel_name: str,
|
||||
platform: str,
|
||||
config_manager: ConfigManager,
|
||||
force: bool = False,
|
||||
autotune_effort: str = "quick",
|
||||
) -> AutotuneResult:
|
||||
logger.debug(
|
||||
"Starting autotune for kernel '%s' with effort='%s'",
|
||||
kernel_name,
|
||||
autotune_effort,
|
||||
)
|
||||
kernel_wrapper = get_kernel_by_name(kernel_name)
|
||||
if kernel_wrapper is None:
|
||||
error_msg = f"Kernel '{kernel_name}' not found in registry"
|
||||
logger.error(error_msg)
|
||||
return AutotuneResult(
|
||||
status="error",
|
||||
message=error_msg,
|
||||
successful=0,
|
||||
failed=0,
|
||||
configs={},
|
||||
)
|
||||
|
||||
try:
|
||||
inputs_dict = kernel_wrapper.get_inputs()
|
||||
except NotImplementedError:
|
||||
error_msg = f"Kernel '{kernel_name}' has no input generator registered"
|
||||
logger.error(error_msg)
|
||||
return AutotuneResult(
|
||||
status="error",
|
||||
message=error_msg,
|
||||
successful=0,
|
||||
failed=0,
|
||||
configs={},
|
||||
)
|
||||
|
||||
try:
|
||||
logger.info(
|
||||
"Autotuning kernel '%s' for platform '%s' with %d configs",
|
||||
kernel_name,
|
||||
platform,
|
||||
len(inputs_dict),
|
||||
)
|
||||
|
||||
configs_to_autotune = {}
|
||||
if not force:
|
||||
existing_configs = config_manager.get_platform_configs(
|
||||
kernel_name, platform
|
||||
)
|
||||
for config_key, inputs in inputs_dict.items():
|
||||
if config_key in existing_configs:
|
||||
logger.debug(
|
||||
"Config '%s' already exists for platform '%s', skipping",
|
||||
config_key,
|
||||
platform,
|
||||
)
|
||||
else:
|
||||
configs_to_autotune[config_key] = inputs
|
||||
else:
|
||||
logger.debug("Force mode enabled, will re-autotune all configs")
|
||||
configs_to_autotune = inputs_dict
|
||||
|
||||
if not configs_to_autotune:
|
||||
logger.info(
|
||||
"All configs already exist for kernel '%s' on platform '%s'. "
|
||||
"Use --force to re-autotune.",
|
||||
kernel_name,
|
||||
platform,
|
||||
)
|
||||
return AutotuneResult(
|
||||
status="skipped",
|
||||
message="All configs already exist",
|
||||
successful=0,
|
||||
failed=0,
|
||||
configs={},
|
||||
)
|
||||
|
||||
total_start_time = time.time()
|
||||
autotuned_configs = {}
|
||||
failed_configs = []
|
||||
|
||||
for config_key, inputs in configs_to_autotune.items():
|
||||
logger.info("Autotuning config: %s", config_key)
|
||||
logger.debug(
|
||||
"Input shapes: %s",
|
||||
[getattr(inp, "shape", type(inp).__name__) for inp in inputs],
|
||||
)
|
||||
|
||||
try:
|
||||
config_start_time = time.time()
|
||||
config = kernel_wrapper.run_autotune(inputs, autotune_effort)
|
||||
config_duration = time.time() - config_start_time
|
||||
|
||||
# Save immediately for checkpointing
|
||||
config_manager.save_configs(kernel_name, platform, {config_key: config})
|
||||
|
||||
autotuned_configs[config_key] = config
|
||||
logger.debug("Config details: %s", config)
|
||||
|
||||
logger.info(
|
||||
"✓ Autotuned and saved config '%s' (%.2fs)",
|
||||
config_key,
|
||||
config_duration,
|
||||
)
|
||||
|
||||
except (RuntimeError, ValueError, OSError) as e:
|
||||
logger.exception(
|
||||
"Failed to autotune config '%s': %s",
|
||||
config_key,
|
||||
e,
|
||||
)
|
||||
failed_configs.append(config_key)
|
||||
|
||||
total_duration = time.time() - total_start_time
|
||||
successful = len(autotuned_configs)
|
||||
failed = len(failed_configs)
|
||||
|
||||
logger.info(
|
||||
"Completed autotuning for kernel '%s': %d successful, %d failed (%.2fs)",
|
||||
kernel_name,
|
||||
successful,
|
||||
failed,
|
||||
total_duration,
|
||||
)
|
||||
|
||||
status = "success" if failed == 0 else "partial"
|
||||
return AutotuneResult(
|
||||
status=status,
|
||||
successful=successful,
|
||||
failed=failed,
|
||||
configs=autotuned_configs,
|
||||
)
|
||||
|
||||
except (KeyError, RuntimeError, ValueError, OSError) as e:
|
||||
error_msg = f"Unexpected error: {e}"
|
||||
logger.exception("Failed to autotune kernel '%s': %s", kernel_name, e)
|
||||
return AutotuneResult(
|
||||
status="error",
|
||||
message=error_msg,
|
||||
successful=0,
|
||||
failed=0,
|
||||
configs={},
|
||||
)
|
||||
|
||||
|
||||
def summarize_results(results: dict[str, AutotuneResult]) -> bool:
|
||||
logger.info("=" * 50)
|
||||
logger.info("Autotuning Results Summary")
|
||||
logger.info("=" * 50)
|
||||
|
||||
total_successful = 0
|
||||
total_failed = 0
|
||||
success_kernels = []
|
||||
partial_kernels = []
|
||||
error_kernels = []
|
||||
skipped_kernels = []
|
||||
|
||||
for kernel_name, result in results.items():
|
||||
total_successful += result.successful
|
||||
total_failed += result.failed
|
||||
|
||||
if result.status == "success":
|
||||
success_kernels.append(f"{kernel_name} ({result.successful} configs)")
|
||||
logger.info("✓ %s: %d configs successful", kernel_name, result.successful)
|
||||
elif result.status == "partial":
|
||||
partial_kernels.append(
|
||||
f"{kernel_name} ({result.successful} ok, {result.failed} failed)"
|
||||
)
|
||||
logger.warning(
|
||||
"⚠ %s: %d successful, %d failed",
|
||||
kernel_name,
|
||||
result.successful,
|
||||
result.failed,
|
||||
)
|
||||
elif result.status == "error":
|
||||
error_kernels.append(f"{kernel_name}: {result.message or 'Unknown error'}")
|
||||
logger.error("✗ %s: %s", kernel_name, result.message or "Unknown error")
|
||||
elif result.status == "skipped":
|
||||
skipped_kernels.append(f"{kernel_name}: {result.message or 'Skipped'}")
|
||||
logger.info("- %s: %s", kernel_name, result.message or "Skipped")
|
||||
|
||||
logger.info("=" * 50)
|
||||
logger.info(
|
||||
"Summary: %d total configs (%d successful, %d failed)",
|
||||
total_successful + total_failed,
|
||||
total_successful,
|
||||
total_failed,
|
||||
)
|
||||
logger.info(
|
||||
"Kernels: %d success, %d partial, %d error, %d skipped",
|
||||
len(success_kernels),
|
||||
len(partial_kernels),
|
||||
len(error_kernels),
|
||||
len(skipped_kernels),
|
||||
)
|
||||
|
||||
has_failures = bool(error_kernels or partial_kernels)
|
||||
|
||||
if not has_failures:
|
||||
if total_successful > 0:
|
||||
logger.info("All configs autotuned successfully!")
|
||||
else:
|
||||
logger.info("No new configs were generated (all may already exist)")
|
||||
|
||||
return not has_failures
|
||||
|
||||
|
||||
def get_kernels_to_autotune(requested_kernels: list[str] | None) -> list[str]:
|
||||
all_kernels = get_registered_kernels()
|
||||
if not all_kernels:
|
||||
logger.error("No Helion kernels found in registry")
|
||||
sys.exit(1)
|
||||
|
||||
if not requested_kernels:
|
||||
return list(all_kernels.keys())
|
||||
|
||||
if len(requested_kernels) != len(set(requested_kernels)):
|
||||
duplicates = [
|
||||
k for k in set(requested_kernels) if requested_kernels.count(k) > 1
|
||||
]
|
||||
logger.error("Duplicate kernel names in --kernels flag: %s", duplicates)
|
||||
sys.exit(1)
|
||||
|
||||
kernels_to_autotune = []
|
||||
missing_kernels = []
|
||||
|
||||
for kernel_name in requested_kernels:
|
||||
if kernel_name in all_kernels:
|
||||
kernels_to_autotune.append(kernel_name)
|
||||
else:
|
||||
missing_kernels.append(kernel_name)
|
||||
|
||||
if missing_kernels:
|
||||
logger.error("Kernel(s) not found: %s", missing_kernels)
|
||||
logger.error("Available kernels: %s", list(all_kernels.keys()))
|
||||
sys.exit(1)
|
||||
|
||||
return kernels_to_autotune
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Autotune Helion kernels",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__.split("Usage:")[1] if "Usage:" in __doc__ else "",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--kernels",
|
||||
nargs="+",
|
||||
help="Kernel(s) to autotune (default: all kernels)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--config-dir",
|
||||
type=str,
|
||||
help="Config directory for config files (default: vLLM helion configs dir)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--list",
|
||||
action="store_true",
|
||||
help="List available Helion kernels and exit",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--force",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Force re-autotuning even if configs already exist for the "
|
||||
"platform and config keys"
|
||||
),
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--autotune-effort",
|
||||
type=str,
|
||||
default="quick",
|
||||
help=(
|
||||
"Helion autotune effort level: 'quick' (smaller search) or "
|
||||
"'full' (full search budget) (default: quick)"
|
||||
),
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--verbose",
|
||||
action="store_true",
|
||||
help="Enable verbose logging",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
import logging
|
||||
|
||||
if args.verbose:
|
||||
logging.getLogger("vllm").setLevel(logging.DEBUG)
|
||||
logger.debug("Verbose mode enabled")
|
||||
logger.debug("Arguments: %s", vars(args))
|
||||
else:
|
||||
logging.getLogger("vllm").setLevel(logging.INFO)
|
||||
|
||||
if args.list:
|
||||
list_kernels()
|
||||
return
|
||||
|
||||
if not check_requirements():
|
||||
sys.exit(1)
|
||||
|
||||
platform = get_canonical_gpu_name()
|
||||
logger.info("Detected GPU platform: %s", platform)
|
||||
|
||||
config_manager = (
|
||||
ConfigManager(args.config_dir) if args.config_dir else ConfigManager()
|
||||
)
|
||||
|
||||
try:
|
||||
config_manager.ensure_base_dir_writable()
|
||||
except OSError as e:
|
||||
logger.error("Failed to access config directory: %s", e)
|
||||
sys.exit(1)
|
||||
|
||||
kernels_to_autotune = get_kernels_to_autotune(args.kernels)
|
||||
|
||||
logger.info(
|
||||
"Will autotune %d kernel(s) for platform '%s': %s",
|
||||
len(kernels_to_autotune),
|
||||
platform,
|
||||
kernels_to_autotune,
|
||||
)
|
||||
|
||||
results = {}
|
||||
for kernel_name in kernels_to_autotune:
|
||||
result = autotune_kernel(
|
||||
kernel_name, platform, config_manager, args.force, args.autotune_effort
|
||||
)
|
||||
results[kernel_name] = result
|
||||
|
||||
success = summarize_results(results)
|
||||
sys.exit(0 if success else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,10 +1,29 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from ..utils import compare_two_settings
|
||||
|
||||
|
||||
def test_cpu_offload():
|
||||
@pytest.mark.parametrize("disable_pin_memory", [False, True])
|
||||
@pytest.mark.parametrize("disable_uva", [False, True])
|
||||
def test_cpu_offload(disable_pin_memory, disable_uva):
|
||||
env_vars = {
|
||||
"VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY": str(int(disable_pin_memory)),
|
||||
"VLLM_WEIGHT_OFFLOADING_DISABLE_UVA": str(int(disable_uva)),
|
||||
}
|
||||
|
||||
args = ["--cpu-offload-gb", "1"]
|
||||
|
||||
# cuda graph only works with UVA offloading
|
||||
if disable_uva:
|
||||
args.append("--enforce-eager")
|
||||
|
||||
compare_two_settings(
|
||||
"hmellor/tiny-random-LlamaForCausalLM", [], ["--cpu-offload-gb", "1"]
|
||||
model="hmellor/tiny-random-LlamaForCausalLM",
|
||||
arg1=[],
|
||||
arg2=args,
|
||||
env1=None,
|
||||
env2=env_vars,
|
||||
)
|
||||
|
||||
@@ -229,7 +229,7 @@ def _compare_sp(
|
||||
if chunked_prefill:
|
||||
common_args.append("--enable-chunked-prefill")
|
||||
if eager_mode:
|
||||
common_args.append("--enforce-eager")
|
||||
common_args.append("-cc.cudagraph_mode=none")
|
||||
if runner != "auto":
|
||||
common_args.extend(["--runner", runner])
|
||||
if trust_remote_code:
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config.model import ModelConfig
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
@@ -23,3 +24,20 @@ def test_mm_encoder_attn_backend_hash_updates():
|
||||
mm_encoder_attn_backend=AttentionBackendEnum.FLASH_ATTN
|
||||
).compute_hash()
|
||||
assert base_hash != overridden_hash
|
||||
|
||||
|
||||
def test_language_model_only_does_not_affect_mm_hash():
|
||||
"""language_model_only does not affect the ViT computation graph,
|
||||
so it should not change the multimodal config hash."""
|
||||
base_hash = MultiModalConfig().compute_hash()
|
||||
lm_only_hash = MultiModalConfig(language_model_only=True).compute_hash()
|
||||
assert base_hash == lm_only_hash
|
||||
|
||||
|
||||
def test_language_model_only_affects_model_hash():
|
||||
"""language_model_only affects the LM computation graph,
|
||||
so it should change the model config hash."""
|
||||
model = "llava-hf/llava-1.5-7b-hf"
|
||||
base_hash = ModelConfig(model).compute_hash()
|
||||
lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash()
|
||||
assert base_hash != lm_only_hash
|
||||
|
||||
@@ -419,7 +419,6 @@ class HfRunner:
|
||||
self.tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast" = (
|
||||
AutoTokenizer.from_pretrained(
|
||||
model_name,
|
||||
dtype=dtype,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
)
|
||||
@@ -430,7 +429,6 @@ class HfRunner:
|
||||
|
||||
self.processor = AutoProcessor.from_pretrained(
|
||||
model_name,
|
||||
dtype=dtype,
|
||||
trust_remote_code=trust_remote_code,
|
||||
)
|
||||
if skip_tokenizer_init:
|
||||
|
||||
@@ -39,7 +39,6 @@ def test_min_tokens_with_stop(min_tokens: int, stop: str, truth: str):
|
||||
mm_features=None,
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
eos_token_id=None,
|
||||
arrival_time=0.0,
|
||||
lora_request=None,
|
||||
cache_salt=None,
|
||||
|
||||
@@ -35,7 +35,6 @@ def _make_request(stop, include_stop_str_in_output: bool, min_tokens: int = 0):
|
||||
mm_features=None,
|
||||
sampling_params=params,
|
||||
pooling_params=None,
|
||||
eos_token_id=None,
|
||||
arrival_time=0.0,
|
||||
lora_request=None,
|
||||
cache_salt=None,
|
||||
|
||||
@@ -19,6 +19,8 @@ from vllm.distributed import (
|
||||
tensor_model_parallel_all_reduce,
|
||||
tensor_model_parallel_reduce_scatter,
|
||||
)
|
||||
from vllm.distributed.parallel_state import GroupCoordinator, TensorMetadata
|
||||
from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
|
||||
|
||||
from ..utils import (
|
||||
init_test_distributed_environment,
|
||||
@@ -200,6 +202,111 @@ def send_recv_tensor_dict_test_worker(
|
||||
torch.testing.assert_close(recv_dict["f"], test_dict["f"])
|
||||
|
||||
|
||||
class _DummyWork:
|
||||
def __init__(self) -> None:
|
||||
self.wait_calls = 0
|
||||
|
||||
def wait(self) -> None:
|
||||
self.wait_calls += 1
|
||||
|
||||
|
||||
class _DummyAllGatherGroup:
|
||||
def __init__(self, world_size: int, rank_in_group: int) -> None:
|
||||
self.world_size = world_size
|
||||
self.rank_in_group = rank_in_group
|
||||
|
||||
def all_gather(self, t: torch.Tensor, dim: int = 0) -> torch.Tensor:
|
||||
# duplicate local slice across ranks.
|
||||
assert dim == 0
|
||||
return torch.cat([t for _ in range(self.world_size)], dim=0)
|
||||
|
||||
|
||||
def _make_group_for_unit_test(
|
||||
rank_in_group: int = 0, world_size: int = 2
|
||||
) -> GroupCoordinator:
|
||||
# avoid running GroupCoordinator.__init__ (it wires up real process groups).
|
||||
g = GroupCoordinator.__new__(GroupCoordinator)
|
||||
g.world_size = world_size
|
||||
g.rank_in_group = rank_in_group
|
||||
g.ranks = list(range(world_size))
|
||||
g.use_cpu_custom_send_recv = False
|
||||
g.device_group = None
|
||||
g.cpu_group = None
|
||||
return g
|
||||
|
||||
|
||||
def test_irecv_tensor_dict_send_allgather_postprocess_binds_keys(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
def fake_irecv(t: torch.Tensor, *args: Any, **kwargs: Any) -> _DummyWork:
|
||||
t.fill_(1)
|
||||
return _DummyWork()
|
||||
|
||||
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
|
||||
monkeypatch.setattr(torch.distributed, "irecv", fake_irecv)
|
||||
|
||||
g = _make_group_for_unit_test(rank_in_group=0, world_size=2)
|
||||
# 2 tensors so we can catch late-binding bugs in postprocess closures.
|
||||
metadata_list = [
|
||||
("a", TensorMetadata("cpu", torch.int32, torch.Size([4]))),
|
||||
("b", TensorMetadata("cpu", torch.int32, torch.Size([4]))),
|
||||
]
|
||||
g.recv_object = lambda src=None: metadata_list # type: ignore[method-assign]
|
||||
|
||||
ag = _DummyAllGatherGroup(world_size=2, rank_in_group=0)
|
||||
td, handles, postprocess = g.irecv_tensor_dict(all_gather_group=ag)
|
||||
|
||||
assert td is not None
|
||||
assert len(handles) == 2
|
||||
assert len(postprocess) == 2
|
||||
|
||||
# before postprocess, dict holds the TP slice (shape 2).
|
||||
assert td["a"].shape == torch.Size([2])
|
||||
assert td["b"].shape == torch.Size([2])
|
||||
|
||||
# simulate worker-side "defer wait": wait + postprocess later.
|
||||
for handle in handles:
|
||||
handle.wait()
|
||||
for fn in postprocess:
|
||||
fn()
|
||||
|
||||
# after postprocess, dict values are reconstructed to full shape (shape 4),
|
||||
# and each key should be updated independently
|
||||
assert td["a"].shape == torch.Size([4])
|
||||
assert td["b"].shape == torch.Size([4])
|
||||
torch.testing.assert_close(td["a"], torch.ones(4, dtype=torch.int32))
|
||||
torch.testing.assert_close(td["b"], torch.ones(4, dtype=torch.int32))
|
||||
|
||||
|
||||
def test_async_intermediate_tensors_lazy_wait() -> None:
|
||||
work = _DummyWork()
|
||||
post_calls = {"n": 0}
|
||||
|
||||
def post() -> None:
|
||||
post_calls["n"] += 1
|
||||
|
||||
it = AsyncIntermediateTensors(
|
||||
{"x": torch.tensor([1])},
|
||||
comm_handles=[work],
|
||||
comm_postprocess=[post],
|
||||
)
|
||||
|
||||
# accessing non-tensor attributes should not trigger wait.
|
||||
assert it.kv_connector_output is None
|
||||
assert work.wait_calls == 0
|
||||
assert post_calls["n"] == 0
|
||||
|
||||
# first access of `.tensors` triggers wait + postprocess.
|
||||
_ = it.tensors
|
||||
assert work.wait_calls == 1
|
||||
assert post_calls["n"] == 1
|
||||
|
||||
# subsequent access should not re-wait.
|
||||
_ = it.tensors
|
||||
assert work.wait_calls == 1
|
||||
assert post_calls["n"] == 1
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def send_recv_test_worker(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
||||
@@ -145,6 +145,7 @@ async def test_request_cancellation(server: RemoteOpenAIServer):
|
||||
model=MODEL_NAME,
|
||||
max_tokens=10000,
|
||||
extra_body={"min_tokens": 10000},
|
||||
temperature=0.0,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
@@ -163,7 +164,7 @@ async def test_request_cancellation(server: RemoteOpenAIServer):
|
||||
# be able to respond to this one within the timeout
|
||||
client = server.get_async_client(timeout=5)
|
||||
response = await client.chat.completions.create(
|
||||
messages=chat_input, model=MODEL_NAME, max_tokens=10
|
||||
messages=chat_input, model=MODEL_NAME, max_tokens=10, temperature=0.0
|
||||
)
|
||||
|
||||
assert len(response.choices) == 1
|
||||
|
||||
@@ -17,6 +17,7 @@ from transformers import AutoTokenizer
|
||||
from tests.conftest import LocalAssetServer
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm import version
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
|
||||
MODELS = {
|
||||
"text": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
|
||||
@@ -315,14 +316,26 @@ async def test_abort_metrics_reset(
|
||||
client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt_ids,
|
||||
max_tokens=100, # Long generation to give time to abort
|
||||
max_tokens=500, # Long generation to give time to abort
|
||||
temperature=0.0,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
# Wait a bit for requests to start processing
|
||||
await asyncio.sleep(0.5)
|
||||
# Poll until we see running requests rather than using a fixed sleep,
|
||||
# since generation speed varies across hardware.
|
||||
try:
|
||||
await _poll_until(
|
||||
lambda: _get_running_metrics_from_api(server)[0] > 0,
|
||||
timeout=10.0,
|
||||
interval=0.1,
|
||||
description="running_requests > 0",
|
||||
)
|
||||
except TimeoutError:
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
pytest.fail("Requests never appeared as running in metrics")
|
||||
|
||||
# Check that we have running requests
|
||||
running_requests, waiting_requests, kv_cache_usage = _get_running_metrics_from_api(
|
||||
@@ -336,13 +349,15 @@ async def test_abort_metrics_reset(
|
||||
# Cancel all tasks to abort the requests
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
# Wait for cancellations to be processed
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
# Check that metrics have reset to zero
|
||||
response = requests.get(server.url_for("metrics"))
|
||||
assert response.status_code == HTTPStatus.OK
|
||||
# Poll until metrics reset rather than using a fixed sleep
|
||||
await _poll_until(
|
||||
lambda: _get_running_metrics_from_api(server) == (0, 0, 0),
|
||||
timeout=10.0,
|
||||
interval=0.2,
|
||||
description="gauge metrics back to zero",
|
||||
)
|
||||
|
||||
# Verify running and waiting requests counts and KV cache usage are zero
|
||||
running_requests_after, waiting_requests_after, kv_cache_usage_after = (
|
||||
@@ -360,6 +375,18 @@ async def test_abort_metrics_reset(
|
||||
)
|
||||
|
||||
|
||||
async def _poll_until(
|
||||
predicate, *, timeout: float, interval: float = 0.5, description: str = "condition"
|
||||
):
|
||||
"""Poll until predicate() returns True, or raise TimeoutError."""
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
if predicate():
|
||||
return
|
||||
await asyncio.sleep(interval)
|
||||
raise TimeoutError(f"Timed out after {timeout}s waiting for: {description}")
|
||||
|
||||
|
||||
def _get_running_metrics_from_api(server: RemoteOpenAIServer):
|
||||
"""Return (running_count, waiting_count, kv_cache_usage)"""
|
||||
|
||||
@@ -399,7 +426,7 @@ def test_metrics_exist_run_batch():
|
||||
input_batch = """{"custom_id": "request-0", "method": "POST", "url": "/v1/embeddings", "body": {"model": "intfloat/multilingual-e5-small", "input": "You are a helpful assistant."}}""" # noqa: E501
|
||||
|
||||
base_url = "0.0.0.0"
|
||||
port = "8001"
|
||||
port = str(get_open_port())
|
||||
server_url = f"http://{base_url}:{port}"
|
||||
|
||||
with (
|
||||
@@ -427,17 +454,32 @@ def test_metrics_exist_run_batch():
|
||||
],
|
||||
)
|
||||
|
||||
def is_server_up(url):
|
||||
try:
|
||||
|
||||
def is_server_up(url):
|
||||
try:
|
||||
response = requests.get(url)
|
||||
return response.status_code == 200
|
||||
except requests.ConnectionError:
|
||||
return False
|
||||
|
||||
start = time.time()
|
||||
timeout = 120
|
||||
while not is_server_up(server_url):
|
||||
if proc.poll() is not None:
|
||||
pytest.fail(
|
||||
f"Batch process exited early with returncode={proc.returncode}"
|
||||
)
|
||||
if time.time() - start > timeout:
|
||||
pytest.fail("Batch server did not start within timeout")
|
||||
time.sleep(1)
|
||||
|
||||
response = requests.get(server_url + "/metrics")
|
||||
assert response.status_code == HTTPStatus.OK
|
||||
finally:
|
||||
proc.terminate()
|
||||
try:
|
||||
response = requests.get(url)
|
||||
return response.status_code == 200
|
||||
except requests.ConnectionError:
|
||||
return False
|
||||
|
||||
while not is_server_up(server_url):
|
||||
time.sleep(1)
|
||||
|
||||
response = requests.get(server_url + "/metrics")
|
||||
assert response.status_code == HTTPStatus.OK
|
||||
|
||||
proc.wait()
|
||||
proc.wait(timeout=15)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.wait(timeout=5)
|
||||
|
||||
@@ -195,18 +195,15 @@ def test_chat_batch_failure_cleanup(llm_for_failure_test):
|
||||
valid_msg = [{"role": "user", "content": "Hello"}]
|
||||
long_text = "This is a very long text to test the error " * 50
|
||||
invalid_msg = [{"role": "user", "content": long_text}]
|
||||
batch_1 = [
|
||||
valid_msg,
|
||||
valid_msg,
|
||||
invalid_msg,
|
||||
]
|
||||
batch_2 = [
|
||||
valid_msg,
|
||||
valid_msg,
|
||||
]
|
||||
|
||||
batch_1 = [valid_msg, valid_msg, invalid_msg]
|
||||
batch_2 = [valid_msg, valid_msg]
|
||||
sampling_params = SamplingParams(temperature=0, max_tokens=10)
|
||||
|
||||
with pytest.raises(ValueError, match="context length is only"):
|
||||
llm.chat(batch_1, sampling_params=sampling_params)
|
||||
assert llm.llm_engine.get_num_unfinished_requests() == 0
|
||||
|
||||
outputs_2 = llm.chat(batch_2, sampling_params=sampling_params)
|
||||
assert len(outputs_2) == len(batch_2)
|
||||
assert llm.llm_engine.get_num_unfinished_requests() == 0
|
||||
|
||||
@@ -6,7 +6,6 @@ import time
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import requests
|
||||
from openai import BadRequestError, NotFoundError, OpenAI
|
||||
from openai_harmony import (
|
||||
Message,
|
||||
@@ -513,11 +512,9 @@ async def test_code_interpreter(client: OpenAI, model_name: str):
|
||||
|
||||
|
||||
def get_weather(latitude, longitude):
|
||||
response = requests.get(
|
||||
f"https://api.open-meteo.com/v1/forecast?latitude={latitude}&longitude={longitude}¤t=temperature_2m,wind_speed_10m&hourly=temperature_2m,relative_humidity_2m,wind_speed_10m" # noqa
|
||||
)
|
||||
data = response.json()
|
||||
return data["current"]["temperature_2m"]
|
||||
# Return a static temperature value to avoid flaky SSL/network errors
|
||||
# from calling the external api.open-meteo.com API in CI.
|
||||
return 15.0
|
||||
|
||||
|
||||
def get_place_to_travel():
|
||||
|
||||
@@ -172,19 +172,26 @@ async def test_mcp_tool_call(client: OpenAI, model_name: str):
|
||||
|
||||
assert response is not None
|
||||
assert response.status == "completed"
|
||||
assert response.output[0].type == "reasoning"
|
||||
assert response.output[1].type == "mcp_call"
|
||||
assert type(response.output[1].arguments) is str
|
||||
assert type(response.output[1].output) is str
|
||||
assert response.output[2].type == "reasoning"
|
||||
# make sure the correct math is in the final output
|
||||
assert response.output[3].type == "message"
|
||||
assert any(s in response.output[3].content[0].text for s in ("56088", "56,088"))
|
||||
|
||||
# test raw input_messages / output_messages
|
||||
# The model may produce multiple reasoning/mcp_call rounds before the
|
||||
# final message, so validate structurally rather than by exact index.
|
||||
output_types = [o.type for o in response.output]
|
||||
assert "reasoning" in output_types
|
||||
mcp_calls = [o for o in response.output if o.type == "mcp_call"]
|
||||
assert len(mcp_calls) >= 1
|
||||
assert type(mcp_calls[0].arguments) is str
|
||||
assert type(mcp_calls[0].output) is str
|
||||
|
||||
# The final output should be a message containing the correct answer
|
||||
assert response.output[-1].type == "message"
|
||||
assert any(s in response.output[-1].content[0].text for s in ("56088", "56,088"))
|
||||
|
||||
# Test raw input_messages / output_messages
|
||||
assert len(response.input_messages) == 1
|
||||
assert len(response.output_messages) == 3
|
||||
assert any(s in response.output_messages[2]["message"] for s in ("56088", "56,088"))
|
||||
assert len(response.output_messages) >= 3
|
||||
assert any(
|
||||
s in response.output_messages[-1]["message"] for s in ("56088", "56,088")
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -4,8 +4,17 @@
|
||||
"""Unit tests for ResponsesRequest.to_sampling_params() parameter mapping."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from openai.types.responses.response_format_text_json_schema_config import (
|
||||
ResponseFormatTextJSONSchemaConfig,
|
||||
)
|
||||
from pydantic import ValidationError
|
||||
|
||||
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
|
||||
from vllm.entrypoints.openai.responses.protocol import (
|
||||
ResponsesRequest,
|
||||
ResponseTextConfig,
|
||||
)
|
||||
from vllm.sampling_params import StructuredOutputsParams
|
||||
|
||||
|
||||
class TestResponsesRequestSamplingParams:
|
||||
@@ -76,9 +85,6 @@ class TestResponsesRequestSamplingParams:
|
||||
|
||||
def test_seed_bounds_validation(self):
|
||||
"""Test that seed values outside torch.long bounds are rejected."""
|
||||
import torch
|
||||
from pydantic import ValidationError
|
||||
|
||||
# Test seed below minimum
|
||||
with pytest.raises(ValidationError) as exc_info:
|
||||
ResponsesRequest(
|
||||
@@ -111,3 +117,40 @@ class TestResponsesRequestSamplingParams:
|
||||
seed=torch.iinfo(torch.long).max,
|
||||
)
|
||||
assert request_max.seed == torch.iinfo(torch.long).max
|
||||
|
||||
def test_structured_outputs_passed_through(self):
|
||||
"""Test that structured_outputs field is passed to SamplingParams."""
|
||||
structured_outputs = StructuredOutputsParams(grammar="root ::= 'hello'")
|
||||
request = ResponsesRequest(
|
||||
model="test-model",
|
||||
input="test input",
|
||||
structured_outputs=structured_outputs,
|
||||
)
|
||||
|
||||
sampling_params = request.to_sampling_params(default_max_tokens=1000)
|
||||
|
||||
assert sampling_params.structured_outputs is not None
|
||||
assert sampling_params.structured_outputs.grammar == "root ::= 'hello'"
|
||||
|
||||
def test_structured_outputs_and_json_schema_conflict(self):
|
||||
"""Test that specifying both structured_outputs and json_schema raises."""
|
||||
structured_outputs = StructuredOutputsParams(grammar="root ::= 'hello'")
|
||||
text_config = ResponseTextConfig()
|
||||
text_config.format = ResponseFormatTextJSONSchemaConfig(
|
||||
type="json_schema",
|
||||
name="test",
|
||||
schema={"type": "object"},
|
||||
)
|
||||
request = ResponsesRequest(
|
||||
model="test-model",
|
||||
input="test input",
|
||||
structured_outputs=structured_outputs,
|
||||
text=text_config,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
request.to_sampling_params(default_max_tokens=1000)
|
||||
|
||||
assert "Cannot specify both structured_outputs and text.format" in str(
|
||||
exc_info.value
|
||||
)
|
||||
|
||||
@@ -54,16 +54,22 @@ class MockModelConfig:
|
||||
media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
|
||||
skip_tokenizer_init = False
|
||||
is_encoder_decoder: bool = False
|
||||
is_multimodal_model: bool = False
|
||||
|
||||
def get_diff_sampling_param(self):
|
||||
return self.diff_sampling_param or {}
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockVllmConfig:
|
||||
model_config: MockModelConfig
|
||||
|
||||
|
||||
def _build_renderer(model_config: MockModelConfig):
|
||||
_, tokenizer_name, _, kwargs = tokenizer_args_from_config(model_config)
|
||||
|
||||
return HfRenderer(
|
||||
model_config,
|
||||
return HfRenderer.from_config(
|
||||
MockVllmConfig(model_config),
|
||||
tokenizer_kwargs={**kwargs, "tokenizer_name": tokenizer_name},
|
||||
)
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user