forked from Karylab-cklius/vllm
Compare commits
18
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e0dd86f6c3 | ||
|
|
01f19dee7f | ||
|
|
2c06cf3486 | ||
|
|
e6f710a87f | ||
|
|
c245d35ff4 | ||
|
|
f8ac0c7cf0 | ||
|
|
ebf862c351 | ||
|
|
8d8062d0a7 | ||
|
|
985961345a | ||
|
|
706a04d34b | ||
|
|
22631f80a0 | ||
|
|
2cc008e7b4 | ||
|
|
5d5c776444 | ||
|
|
592ae6805c | ||
|
|
7b1bc0a3eb | ||
|
|
c0879d9483 | ||
|
|
f5f9878514 | ||
|
|
2ce95a761b |
@@ -121,25 +121,7 @@ steps:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
|
||||
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
|
||||
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
--load \
|
||||
-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.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
# 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) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
@@ -152,25 +134,7 @@ steps:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
|
||||
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
|
||||
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
--load \
|
||||
-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.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 12.9"
|
||||
@@ -180,24 +144,7 @@ steps:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
|
||||
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=12.9.1 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
|
||||
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
--load \
|
||||
-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=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
|
||||
# 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)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
@@ -210,24 +157,7 @@ steps:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
|
||||
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=12.9.1 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
|
||||
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
--load \
|
||||
-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=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
@@ -237,27 +167,7 @@ steps:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
|
||||
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg UBUNTU_VERSION=24.04 \
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
|
||||
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
--load \
|
||||
-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.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
@@ -269,27 +179,7 @@ steps:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
|
||||
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg UBUNTU_VERSION=24.04 \
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
|
||||
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
--load \
|
||||
-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.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
@@ -299,26 +189,7 @@ steps:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
|
||||
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=12.9.1 \
|
||||
--build-arg UBUNTU_VERSION=24.04 \
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
|
||||
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
--load \
|
||||
-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=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
@@ -330,26 +201,7 @@ steps:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 936637512419.dkr.ecr.us-east-1.amazonaws.com || true"
|
||||
- "docker buildx create --name release-builder --driver docker-container --use 2>/dev/null || docker buildx use release-builder"
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=12.9.1 \
|
||||
--build-arg UBUNTU_VERSION=24.04 \
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-postmerge-cache:latest,mode=max \
|
||||
--cache-from type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max \
|
||||
--cache-to type=registry,ref=936637512419.dkr.ecr.us-east-1.amazonaws.com/vllm-ci-test-cache:release-$BUILDKITE_COMMIT,mode=max,compression=zstd \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
--load \
|
||||
-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=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
|
||||
|
||||
- block: "Build release image for x86_64 CPU"
|
||||
|
||||
+14
-14
@@ -388,10 +388,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
@@ -1647,10 +1647,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
@@ -1951,8 +1951,8 @@ steps:
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
#------------------------------------------------------- mi300 · quantization --------------------------------------------------------#
|
||||
|
||||
@@ -2930,10 +2930,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
|
||||
@@ -95,11 +95,13 @@ steps:
|
||||
- tests/kernels/moe/test_deepgemm.py
|
||||
- tests/kernels/moe/test_batched_deepgemm.py
|
||||
- tests/kernels/attention/test_deepgemm_attention.py
|
||||
- tests/quantization/test_cutlass_w4a16.py
|
||||
commands:
|
||||
- pytest -v -s kernels/quantization/test_block_fp8.py
|
||||
- pytest -v -s kernels/moe/test_deepgemm.py
|
||||
- pytest -v -s kernels/moe/test_batched_deepgemm.py
|
||||
- pytest -v -s kernels/attention/test_deepgemm_attention.py
|
||||
- pytest -v -s quantization/test_cutlass_w4a16.py
|
||||
|
||||
- label: Kernels (B200)
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -113,10 +113,10 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# for multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
|
||||
@@ -44,10 +44,10 @@ steps:
|
||||
#- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
|
||||
#- python3 basic/offline_inference/embed.py # TODO
|
||||
# for multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
|
||||
@@ -69,9 +69,9 @@ steps:
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
- label: Transformers Backward Compatibility Models Test
|
||||
working_dir: "/vllm-workspace/"
|
||||
@@ -83,7 +83,7 @@ steps:
|
||||
- pytest -v -s tests/models/test_transformers.py
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/offline_inference/basic/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
+1
-5
@@ -389,11 +389,7 @@ pull_request_rules:
|
||||
- files~=^tests/entrypoints/anthropic/.*tool.*
|
||||
- files~=^vllm/tool_parsers/
|
||||
- files=docs/features/tool_calling.md
|
||||
- files~=^examples/tool_chat_*
|
||||
- files=examples/offline_inference/chat_with_tools.py
|
||||
- files=examples/online_serving/openai_chat_completion_client_with_tools_required.py
|
||||
- files=examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py
|
||||
- files=examples/online_serving/openai_chat_completion_client_with_tools.py
|
||||
- files~=^examples/tool_calling/
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
|
||||
@@ -0,0 +1,324 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Benchmarks FP8 vs BF16 ViT attention via FlashInfer cuDNN backend.
|
||||
#
|
||||
# == Usage Examples ==
|
||||
#
|
||||
# Benchmark mode (default, FlashInfer CUDAGraph Bench)
|
||||
# python3 benchmark_vit_fp8_attn.py
|
||||
#
|
||||
# Profile mode (PyTorch profiler, saves TensorBoard traces):
|
||||
# python3 benchmark_vit_fp8_attn.py --profile
|
||||
# python3 benchmark_vit_fp8_attn.py --profile --profile-output-dir ./profile_traces
|
||||
#
|
||||
# Custom seq_lens:
|
||||
# python3 benchmark_vit_fp8_attn.py --seq-lens 4096 8192 16384
|
||||
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.profiler import ProfilerActivity, profile, record_function
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
# Qwen3-VL defaults
|
||||
NUM_HEADS = 16
|
||||
HEAD_DIM = 72
|
||||
DEFAULT_SEQ_LENS = [2304, 4096, 8192, 16384]
|
||||
|
||||
|
||||
def _setup_fp8_attention(num_heads: int, head_dim: int) -> tuple:
|
||||
"""Create FP8 and BF16 attention modules + workspace."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
_get_flashinfer_workspace_buffer,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
old_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
backend_patch = patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
)
|
||||
|
||||
# FP8 attention
|
||||
mm_config_fp8 = MultiModalConfig(mm_encoder_attn_dtype="fp8")
|
||||
vllm_config_fp8 = VllmConfig()
|
||||
vllm_config_fp8.model_config = SimpleNamespace(multimodal_config=mm_config_fp8)
|
||||
with set_current_vllm_config(vllm_config_fp8), backend_patch:
|
||||
attn_fp8 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
# BF16 attention (no FP8)
|
||||
with set_current_vllm_config(VllmConfig()), backend_patch:
|
||||
attn_bf16 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
torch.set_default_dtype(old_dtype)
|
||||
|
||||
workspace = _get_flashinfer_workspace_buffer()
|
||||
return attn_fp8, attn_bf16, workspace
|
||||
|
||||
|
||||
def _build_meta(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
fp8: bool,
|
||||
):
|
||||
"""Build cu_seqlens, max_seqlen, sequence_lengths."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
cu_np = np.array([0, seq_len], dtype=np.int32)
|
||||
fp8_padded = num_heads * round_up(head_dim, 16) if fp8 else None
|
||||
|
||||
seq_lengths = MMEncoderAttention.maybe_compute_seq_lens(
|
||||
AttentionBackendEnum.FLASHINFER, cu_np, torch.device("cuda")
|
||||
)
|
||||
max_seqlen = torch.tensor(
|
||||
MMEncoderAttention.compute_max_seqlen(AttentionBackendEnum.FLASHINFER, cu_np),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_np,
|
||||
num_heads * head_dim,
|
||||
1,
|
||||
torch.device("cuda"),
|
||||
fp8_padded_hidden_size=fp8_padded,
|
||||
)
|
||||
return cu_seqlens, max_seqlen, seq_lengths
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
seq_lens: list[int],
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
method: str,
|
||||
):
|
||||
"""Benchmark FP8 vs BF16 attention across seq_lens.
|
||||
|
||||
Uses FlashInfer GPU-level timing to measure pure kernel time,
|
||||
excluding CPU launch overhead.
|
||||
"""
|
||||
if method == "cupti":
|
||||
from flashinfer.testing import bench_gpu_time_with_cupti as bench_fn
|
||||
|
||||
bench_fn = partial(bench_fn, use_cuda_graph=True, cold_l2_cache=False)
|
||||
elif method == "cudagraph":
|
||||
from flashinfer.testing import (
|
||||
bench_gpu_time_with_cudagraph as bench_fn,
|
||||
)
|
||||
|
||||
bench_fn = partial(bench_fn, cold_l2_cache=False)
|
||||
else:
|
||||
raise ValueError(f"Invalid method: {method}")
|
||||
|
||||
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
|
||||
|
||||
print(f"Timing method: {method}")
|
||||
print(f"{'seq_len':>8} {'BF16 (us)':>12} {'FP8 (us)':>12} {'Speedup':>10}")
|
||||
print("-" * 46)
|
||||
|
||||
for seq_len in seq_lens:
|
||||
torch.manual_seed(42)
|
||||
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
|
||||
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
|
||||
cu_bf16 = cu_fp8.clone()
|
||||
|
||||
def bf16_fn(q=q, k=k, v=v, cu=cu_bf16, ms=max_s, sl=seq_l):
|
||||
attn_bf16._forward_flashinfer(q, k, v, cu, ms, sl)
|
||||
|
||||
def fp8_fn(q=q, k=k, v=v, cu=cu_fp8, ms=max_s, sl=seq_l):
|
||||
attn_fp8._forward_flashinfer(q, k, v, cu, ms, sl)
|
||||
|
||||
# bench_fn returns List[float] of per-iteration times in ms
|
||||
bf16_times = bench_fn(bf16_fn)
|
||||
fp8_times = bench_fn(fp8_fn)
|
||||
|
||||
bf16_us = np.median(bf16_times) * 1e3 # ms -> us
|
||||
fp8_us = np.median(fp8_times) * 1e3
|
||||
speedup = bf16_us / fp8_us if fp8_us > 0 else float("inf")
|
||||
|
||||
print(f"{seq_len:>8} {bf16_us:>12.1f} {fp8_us:>12.1f} {speedup:>9.2f}x")
|
||||
|
||||
|
||||
def _make_trace_handler(output_dir: str, worker_name: str, label: str):
|
||||
"""Create a trace handler that saves to TensorBoard and prints summary."""
|
||||
|
||||
def handler(prof):
|
||||
torch.profiler.tensorboard_trace_handler(output_dir, worker_name)(prof)
|
||||
print(f"\n{'=' * 80}")
|
||||
print(label)
|
||||
print(f"{'=' * 80}")
|
||||
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def run_profile(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
warmup: int,
|
||||
output_dir: str,
|
||||
):
|
||||
"""Profile FP8 vs BF16 attention with PyTorch profiler."""
|
||||
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
|
||||
|
||||
torch.manual_seed(42)
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
|
||||
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
|
||||
cu_bf16 = cu_fp8.clone()
|
||||
|
||||
sched = torch.profiler.schedule(wait=0, warmup=warmup, active=1)
|
||||
|
||||
# Profile BF16 (warmup handled by profiler schedule)
|
||||
with profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
||||
schedule=sched,
|
||||
on_trace_ready=_make_trace_handler(
|
||||
output_dir,
|
||||
f"bf16_h{head_dim}_s{seq_len}",
|
||||
f"BF16 Attention (seq_len={seq_len}, heads={num_heads}, "
|
||||
f"head_dim={head_dim})",
|
||||
),
|
||||
) as prof_bf16:
|
||||
for _ in range(warmup + 1):
|
||||
with record_function("bf16_attention"):
|
||||
attn_bf16._forward_flashinfer(
|
||||
q.clone(), k.clone(), v.clone(), cu_bf16, max_s, seq_l
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
prof_bf16.step()
|
||||
|
||||
# Profile FP8 (warmup handled by profiler schedule)
|
||||
with profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
||||
schedule=sched,
|
||||
on_trace_ready=_make_trace_handler(
|
||||
output_dir,
|
||||
f"fp8_h{head_dim}_s{seq_len}",
|
||||
f"FP8 Attention (seq_len={seq_len}, heads={num_heads}, "
|
||||
f"head_dim={head_dim})",
|
||||
),
|
||||
) as prof_fp8:
|
||||
for _ in range(warmup + 1):
|
||||
with record_function("fp8_attention"):
|
||||
attn_fp8._forward_flashinfer(
|
||||
q.clone(), k.clone(), v.clone(), cu_fp8, max_s, seq_l
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
prof_fp8.step()
|
||||
|
||||
print(f"\nTensorBoard traces saved to: {output_dir}")
|
||||
print(f"View with: tensorboard --logdir={output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = FlexibleArgumentParser(description="Benchmark FP8 vs BF16 ViT attention.")
|
||||
parser.add_argument(
|
||||
"--seq-lens",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=DEFAULT_SEQ_LENS,
|
||||
help="Sequence lengths to benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-heads",
|
||||
type=int,
|
||||
default=NUM_HEADS,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--head-dim",
|
||||
type=int,
|
||||
default=HEAD_DIM,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--method",
|
||||
choices=["cupti", "cudagraph"],
|
||||
default="cudagraph",
|
||||
help="GPU timing method: cupti (CUPTI kernel timing) or "
|
||||
"cudagraph (CUDA graph capture/replay). Default: cudagraph",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Warmup iterations (profile mode only)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile",
|
||||
action="store_true",
|
||||
help="Run PyTorch profiler instead of benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-seq-len",
|
||||
type=int,
|
||||
default=8192,
|
||||
help="Sequence length for profiling (default: 8192)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-output-dir",
|
||||
type=str,
|
||||
default="./profile_traces",
|
||||
help="Output directory for TensorBoard traces (default: ./profile_traces)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.profile:
|
||||
run_profile(
|
||||
args.profile_seq_len,
|
||||
args.num_heads,
|
||||
args.head_dim,
|
||||
args.warmup,
|
||||
args.profile_output_dir,
|
||||
)
|
||||
else:
|
||||
run_benchmark(
|
||||
args.seq_lens,
|
||||
args.num_heads,
|
||||
args.head_dim,
|
||||
args.method,
|
||||
)
|
||||
+4
-59
@@ -82,73 +82,18 @@ void launch_persistent_topk(const torch::Tensor& logits,
|
||||
size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
|
||||
if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;
|
||||
|
||||
// Query occupancy for the instantiation that will actually launch;
|
||||
// overestimating it deadlocks the cooperative barrier.
|
||||
int occupancy = 1;
|
||||
cudaError_t occ_err = cudaSuccess;
|
||||
if (vec_size == 4) {
|
||||
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
&occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
|
||||
smem_size);
|
||||
} else if (vec_size == 2) {
|
||||
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
&occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
|
||||
smem_size);
|
||||
} else {
|
||||
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
&occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
|
||||
smem_size);
|
||||
}
|
||||
TORCH_CHECK(occ_err == cudaSuccess,
|
||||
"persistent_topk occupancy query failed: ",
|
||||
cudaGetErrorString(occ_err));
|
||||
cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
&occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
|
||||
smem_size);
|
||||
if (occupancy < 1) occupancy = 1;
|
||||
|
||||
// The cooperative spin-wait barrier only runs when at least one row hits
|
||||
// the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
|
||||
// early-exit, so oversubscription can't deadlock and headroom is wasted.
|
||||
const bool needs_cooperative =
|
||||
static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
|
||||
|
||||
const uint32_t hw_resident_cap =
|
||||
static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
|
||||
uint32_t max_resident_ctas = hw_resident_cap;
|
||||
if (needs_cooperative) {
|
||||
// Reserve one CTA per SM when occupancy allows; fall back to a single
|
||||
// CTA when occupancy == 1 (the most deadlock-prone case — any straggler
|
||||
// kernel that takes the only slot on one SM hangs the barrier). Never
|
||||
// drop below one full group's worth.
|
||||
uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
|
||||
if (max_resident_ctas >= headroom + ctas_per_group) {
|
||||
max_resident_ctas -= headroom;
|
||||
}
|
||||
}
|
||||
uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
|
||||
uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
|
||||
static_cast<uint32_t>(num_rows));
|
||||
if (num_groups == 0) num_groups = 1;
|
||||
uint32_t total_ctas = num_groups * ctas_per_group;
|
||||
|
||||
// If the cooperative launch wouldn't fit, fall back to FilteredTopK
|
||||
// instead of deadlocking. Only relevant when needs_cooperative.
|
||||
if (needs_cooperative && total_ctas > hw_resident_cap) {
|
||||
TORCH_CHECK(max_smem_per_block >= 128 * 1024,
|
||||
"persistent_topk would oversubscribe and the FilteredTopK "
|
||||
"fallback requires >=128KB smem per block (have ",
|
||||
max_smem_per_block, "). total_ctas=", total_ctas,
|
||||
" > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
|
||||
", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
|
||||
", smem=", smem_size, ").");
|
||||
cudaError_t status =
|
||||
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
|
||||
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
|
||||
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
|
||||
static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
|
||||
stream);
|
||||
TORCH_CHECK(status == cudaSuccess,
|
||||
"FilteredTopK fallback failed: ", cudaGetErrorString(status));
|
||||
return;
|
||||
}
|
||||
|
||||
size_t state_bytes = num_groups * sizeof(P::RadixRowState);
|
||||
TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
|
||||
"workspace too small, need ", state_bytes, " bytes");
|
||||
|
||||
+27
-34
@@ -490,46 +490,41 @@ WORKDIR /vllm-workspace
|
||||
RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
|
||||
echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
|
||||
|
||||
# Install Python from source and system dependencies
|
||||
# Install Python and system dependencies
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
software-properties-common \
|
||||
curl \
|
||||
sudo \
|
||||
ffmpeg \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libgl1 \
|
||||
&& if [ ! -z ${DEADSNAKES_MIRROR_URL} ] ; then \
|
||||
if [ ! -z "${DEADSNAKES_GPGKEY_URL}" ] ; then \
|
||||
mkdir -p -m 0755 /etc/apt/keyrings ; \
|
||||
curl -L ${DEADSNAKES_GPGKEY_URL} | gpg --dearmor > /etc/apt/keyrings/deadsnakes.gpg ; \
|
||||
sudo chmod 644 /etc/apt/keyrings/deadsnakes.gpg ; \
|
||||
echo "deb [signed-by=/etc/apt/keyrings/deadsnakes.gpg] ${DEADSNAKES_MIRROR_URL} $(lsb_release -cs) main" > /etc/apt/sources.list.d/deadsnakes.list ; \
|
||||
fi ; \
|
||||
else \
|
||||
for i in 1 2 3; do \
|
||||
add-apt-repository -y ppa:deadsnakes/ppa && break || \
|
||||
{ echo "Attempt $i failed, retrying in 5s..."; sleep 5; }; \
|
||||
done ; \
|
||||
fi \
|
||||
&& apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
python${PYTHON_VERSION} \
|
||||
python${PYTHON_VERSION}-dev \
|
||||
python${PYTHON_VERSION}-venv \
|
||||
libibverbs-dev \
|
||||
build-essential \
|
||||
libssl-dev \
|
||||
libffi-dev \
|
||||
zlib1g-dev \
|
||||
libbz2-dev \
|
||||
libreadline-dev \
|
||||
libsqlite3-dev \
|
||||
libncurses-dev \
|
||||
liblzma-dev \
|
||||
libgdbm-dev \
|
||||
uuid-dev \
|
||||
tk-dev \
|
||||
&& PYTHON_MAJOR_MINOR=${PYTHON_VERSION} \
|
||||
&& PYTHON_FULL_VERSION=$(curl -s https://www.python.org/ftp/python/ \
|
||||
| grep -oE "${PYTHON_MAJOR_MINOR}\.[0-9]+" \
|
||||
| sort -t. -k3 -n | uniq \
|
||||
| tail -1) \
|
||||
&& echo "Building Python ${PYTHON_FULL_VERSION} from source..." \
|
||||
&& curl -fSL https://www.python.org/ftp/python/${PYTHON_FULL_VERSION}/Python-${PYTHON_FULL_VERSION}.tgz -o /tmp/python.tgz \
|
||||
&& tar -xzf /tmp/python.tgz -C /tmp \
|
||||
&& cd /tmp/Python-${PYTHON_FULL_VERSION} \
|
||||
&& ./configure --enable-optimizations --with-ensurepip=install --prefix=/usr/local \
|
||||
&& make -j$(nproc) \
|
||||
&& make install \
|
||||
&& cd / \
|
||||
&& rm -rf /tmp/python.tgz /tmp/Python-${PYTHON_FULL_VERSION} \
|
||||
&& ln -sf /usr/local/bin/python${PYTHON_MAJOR_MINOR} /usr/bin/python3 \
|
||||
&& ln -sf /usr/local/bin/python${PYTHON_MAJOR_MINOR}-config /usr/bin/python3-config \
|
||||
&& ln -sf /usr/local/bin/pip${PYTHON_MAJOR_MINOR} /usr/bin/pip \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
|
||||
&& update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
|
||||
&& ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
|
||||
&& rm -f /usr/lib/python${PYTHON_VERSION}/EXTERNALLY-MANAGED \
|
||||
&& curl -sS ${GET_PIP_URL} | python${PYTHON_VERSION} \
|
||||
&& python3 --version && python3 -m pip --version
|
||||
|
||||
# Install CUDA development tools for runtime JIT compilation
|
||||
@@ -543,11 +538,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
|
||||
cuda-nvrtc-${CUDA_VERSION_DASH} \
|
||||
cuda-cuobjdump-${CUDA_VERSION_DASH} \
|
||||
libcurand-dev-${CUDA_VERSION_DASH} \
|
||||
libcublas-${CUDA_VERSION_DASH} \
|
||||
libcublas-dev-${CUDA_VERSION_DASH} \
|
||||
# Required by fastsafetensors (fixes #20384)
|
||||
libnuma-dev \
|
||||
# numactl CLI for NUMA binding at runtime
|
||||
numactl && \
|
||||
libnuma-dev && \
|
||||
# Fixes nccl_allocator requiring nccl.h at runtime
|
||||
# https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
|
||||
# NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
|
||||
|
||||
@@ -36,7 +36,7 @@ th {
|
||||
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
|
||||
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
|
||||
!!! info "Table key"
|
||||
1. All types: mxfp4, nvfp4, int4, int8, fp8
|
||||
|
||||
@@ -68,7 +68,7 @@ You can pass a single image to the `'image'` field of the multi-modal dictionary
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
|
||||
|
||||
To substitute multiple images inside the same text prompt, you can pass in a list of images instead:
|
||||
|
||||
@@ -101,7 +101,7 @@ To substitute multiple images inside the same text prompt, you can pass in a lis
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
Full example: [examples/offline_inference/vision_language_multi_image.py](../../examples/offline_inference/vision_language_multi_image.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_multi_image_offline.py](../../examples/generate/multimodal/vision_language_multi_image_offline.py)
|
||||
|
||||
If using the [LLM.chat](../models/generative_models.md#llmchat) method, you can pass images directly in the message content using various formats: image URLs, PIL Image objects, or pre-computed embeddings:
|
||||
|
||||
@@ -287,13 +287,13 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown
|
||||
!!! note
|
||||
'process_vision_info' is only applicable to Qwen2.5-VL and similar models.
|
||||
|
||||
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
|
||||
|
||||
### Audio Inputs
|
||||
|
||||
You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary.
|
||||
|
||||
Full example: [examples/offline_inference/audio_language.py](../../examples/offline_inference/audio_language.py)
|
||||
Full example: [examples/generate/multimodal/audio_language_offline.py](../../examples/generate/multimodal/audio_language_offline.py)
|
||||
|
||||
#### Chunking Long Audio for Transcription
|
||||
|
||||
@@ -674,7 +674,7 @@ Then, you can use the OpenAI client as follows:
|
||||
print("Chat completion output:", chat_response.choices[0].message.content)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! tip
|
||||
Loading from local file paths is also supported on vLLM: You can specify the allowed local media path via `--allowed-local-media-path` when launching the API server/engine,
|
||||
@@ -745,7 +745,7 @@ Then, you can use the OpenAI client as follows:
|
||||
print("Chat completion output from image url:", result)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! note
|
||||
By default, the timeout for fetching videos through HTTP URL is `30` seconds.
|
||||
@@ -958,7 +958,7 @@ Alternatively, you can pass `audio_url`, which is the audio counterpart of `imag
|
||||
print("Chat completion output from audio url:", result)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! note
|
||||
By default, the timeout for fetching audios through HTTP URL is `10` seconds.
|
||||
|
||||
@@ -20,6 +20,7 @@ The following are the supported quantization formats for vLLM:
|
||||
- [AMD Quark](quark.md)
|
||||
- [Quantized KV Cache](quantized_kvcache.md)
|
||||
- [TorchAO](torchao.md)
|
||||
- [FP8 ViT Encoder Attention](fp8_vit_attn.md)
|
||||
|
||||
## Supported Hardware
|
||||
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
# FP8 ViT Encoder Attention
|
||||
|
||||
For visual understanding workloads with large images (e.g. QHD, 4K) and relatively
|
||||
short text prompts/generation, the ViT encoder attention can become a significant
|
||||
bottleneck, especially when the text model is quantized (e.g. NVFP4). vLLM
|
||||
supports optional FP8 quantization for the ViT encoder attention via the
|
||||
FlashInfer cuDNN backend. Q/K/V are quantized on-the-fly to FP8 before the
|
||||
cuDNN attention call.
|
||||
|
||||
!!! note
|
||||
- Currently supports Qwen3-VL family models only (`qwen3_vl`, `qwen3_vl_moe`,
|
||||
`qwen3_5`, `qwen3_5_moe`, and other models using Qwen3 ViT).
|
||||
- Dynamic scaling is not compatible with ViT full CUDA graphs.
|
||||
- Performance gains are mostly visible at QHD/4K resolutions or multi-image
|
||||
requests. Smaller images may see no speedup due to quantization overhead
|
||||
(3 quantization kernel launches + un-padding).
|
||||
- FP8 tensor-core speedup is more pronounced on GB300 than GB200.
|
||||
|
||||
## Requirements
|
||||
|
||||
- FlashInfer cuDNN backend with cuDNN >= 9.17.1.
|
||||
|
||||
## Usage
|
||||
|
||||
Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` together
|
||||
with `--mm-encoder-attn-backend FLASHINFER`:
|
||||
|
||||
```bash
|
||||
vllm serve $MODEL \
|
||||
--mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8
|
||||
```
|
||||
|
||||
By default (no scale file), **dynamic scaling** is used: a 16-entry circular
|
||||
buffer of observed Q/K/V amax values drives per-forward scale updates. This
|
||||
matches BF16 accuracy without any calibration but adds a small per-forward
|
||||
overhead.
|
||||
|
||||
## Calibrate-Once, Reuse Workflow (Recommended)
|
||||
|
||||
For production, calibrate static scales on a representative dataset once and
|
||||
reuse them to avoid the dynamic overhead:
|
||||
|
||||
```bash
|
||||
# Step 1: calibrate and save scales (runs dynamic scaling for 16 passes,
|
||||
# then dumps the learned scales to JSON).
|
||||
vllm bench mm-processor \
|
||||
--model $MODEL --mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8 \
|
||||
--mm-encoder-fp8-scale-save-path /path/to/scales.json \
|
||||
--dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \
|
||||
--num-prompts 100
|
||||
|
||||
# Step 2: serve with static scales (no dynamic overhead).
|
||||
vllm serve $MODEL \
|
||||
--mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8 \
|
||||
--mm-encoder-fp8-scale-path /path/to/scales.json
|
||||
```
|
||||
|
||||
Saved scales are multiplied by `--mm-encoder-fp8-scale-save-margin` (default
|
||||
`1.5`) to leave headroom against activation outliers not present in the
|
||||
calibration set. The default has been validated to generalize across datasets
|
||||
(e.g. VisionArena-Chat calibration maintains BF16 accuracy on ChartQA).
|
||||
|
||||
## Scale File Format
|
||||
|
||||
```json
|
||||
{
|
||||
"visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0},
|
||||
"visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0}
|
||||
}
|
||||
```
|
||||
|
||||
Keys `q_scale` / `k_scale` / `v_scale` are accepted as aliases.
|
||||
|
||||
## Performance
|
||||
|
||||
**Core cuDNN attention kernel** (PyTorch profiler, `cudnn_generated_fort_native_sdpa_sm100_flash_fprop`, head_dim=128, seq_len=8192):
|
||||
|
||||
| Hardware | BF16 | FP8 | Speedup |
|
||||
| -------- | ---- | ---- | ------- |
|
||||
| GB200 | 350 us | 312 us | **1.12x** |
|
||||
| GB300 | 300 us | 211 us | **1.42x** |
|
||||
|
||||
**End-to-end encoder forward time** (Qwen3-VL-30B-A3B-Instruct on GB200, 3 images/request):
|
||||
|
||||
| Resolution | BF16 median | FP8 median | Speedup |
|
||||
| ---------- | ----------- | ---------- | ------- |
|
||||
| HD (720x1280) | 31.77 ms | 36.39 ms | 0.87x |
|
||||
| FullHD (1080x1920) | 57.99 ms | 58.73 ms | ~same |
|
||||
| QHD (1440x2560) | 131.83 ms | 122.30 ms | **1.08x** |
|
||||
| 4K (2160x3840) | 543.44 ms | 460.31 ms | **1.18x** |
|
||||
|
||||
Crossover is around FullHD with 3 images/request. At QHD and above, FP8 wins.
|
||||
|
||||
## Accuracy
|
||||
|
||||
ChartQA, Qwen3-VL-8B-Instruct, 500 samples. FP8 static uses scales calibrated
|
||||
on VisionArena-Chat (with default 1.5x margin):
|
||||
|
||||
| Metric | BF16 | FP8 dynamic | FP8 static |
|
||||
| ------ | ---- | ----------- | ---------- |
|
||||
| relaxed_accuracy | 0.780 | 0.776 | 0.780 |
|
||||
| anywhere_accuracy | 0.806 | 0.816 | 0.814 |
|
||||
| exact_match | 0.584 | 0.582 | 0.578 |
|
||||
|
||||
All three configurations match within statistical noise, confirming that
|
||||
static scales calibrated on one dataset generalize to another.
|
||||
@@ -202,7 +202,7 @@ The reasoning content is also available when both tool calling and the reasoning
|
||||
print(f"Arguments: {tool_call.arguments}")
|
||||
```
|
||||
|
||||
For more examples, please refer to [examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py).
|
||||
For more examples, please refer to [examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py).
|
||||
|
||||
## Server-Level Default Chat Template Kwargs
|
||||
|
||||
|
||||
@@ -292,10 +292,10 @@ Pooling models now support token-wise task.
|
||||
|
||||
### Score task
|
||||
|
||||
`score` task is deprecated and will be removed in v0.20. Please use `classify` instead. Only when a
|
||||
classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
`score` task have has been removed in v0.21, use `classify` instead. Only when a classification model outputs num_labels
|
||||
equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
|
||||
### Pooling multitask support
|
||||
|
||||
Pooling multitask support is deprecated and will be removed in v0.20. When the default pooling task is not what you want,
|
||||
Pooling multitask support has been removed in v0.21. When the default pooling task is not what you want,
|
||||
you need to manually specify it via `PoolerConfig(task=<task>)` offline or `--pooler-config.task <task>` online.
|
||||
|
||||
@@ -439,6 +439,7 @@ th {
|
||||
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
|
||||
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
|
||||
| `MiMoV2ProForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
|
||||
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
|
||||
@@ -590,6 +591,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |
|
||||
| `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
|
||||
| `MiDashengLMModel` | MiDashengLM | T + A<sup>+</sup> | `mispeech/midashenglm-7b` | | ✅︎ |
|
||||
| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
|
||||
| `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
|
||||
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
|
||||
|
||||
@@ -251,7 +251,7 @@ The following extra parameters are supported:
|
||||
Our Responses API is compatible with [OpenAI's Responses API](https://platform.openai.com/docs/api-reference/responses);
|
||||
you can use the [official OpenAI Python client](https://github.com/openai/openai-python) to interact with it.
|
||||
|
||||
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/online_serving/openai_responses_client_with_tools.py)
|
||||
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/tool_calling/openai_responses_client_with_tools.py)
|
||||
|
||||
#### Extra parameters
|
||||
|
||||
@@ -279,7 +279,7 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai
|
||||
!!! note
|
||||
To use the Transcriptions API, please install with extra audio dependencies using `pip install vllm[audio]`.
|
||||
|
||||
Code example: [examples/online_serving/openai_transcription_client.py](../../examples/online_serving/openai_transcription_client.py)
|
||||
Code example: [examples/speech_to_text/openai/openai_transcription_client.py](../../examples/speech_to_text/openai/openai_transcription_client.py)
|
||||
|
||||
NOTE: beam search is currently supported in the transcriptions endpoint for encoder-decoder multimodal models, e.g., whisper, but highly inefficient as work for handling the encoder/decoder cache is actively ongoing. This is an active point of ongoing optimization and will be handled properly in the very near future.
|
||||
|
||||
@@ -397,7 +397,7 @@ Please mind that the popular `openai/whisper-large-v3-turbo` model does not supp
|
||||
!!! note
|
||||
To use the Translation API, please install with extra audio dependencies using `pip install vllm[audio]`.
|
||||
|
||||
Code example: [examples/online_serving/openai_translation_client.py](../../examples/online_serving/openai_translation_client.py)
|
||||
Code example: [examples/speech_to_text/openai/openai_translation_client.py](../../examples/speech_to_text/openai/openai_translation_client.py)
|
||||
|
||||
#### Extra Parameters
|
||||
|
||||
|
||||
Executable → Regular
+6
-6
@@ -6,15 +6,15 @@ This folder provides several example scripts on how to inference Qwen2.5-Omni of
|
||||
|
||||
```bash
|
||||
# Audio + image + video
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
-q mixed_modalities
|
||||
|
||||
# Read vision and audio inputs from a single video file
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
-q use_audio_in_video
|
||||
|
||||
# Multiple audios
|
||||
python examples/offline_inference/qwen2_5_omni/only_thinker.py \
|
||||
python examples/generate/multimodal/qwen2_5_omni/only_thinker.py \
|
||||
-q multi_audios
|
||||
```
|
||||
|
||||
@@ -24,16 +24,16 @@ You can also test Qwen2.5-Omni on a single modality:
|
||||
|
||||
```bash
|
||||
# Process audio inputs
|
||||
python examples/offline_inference/audio_language.py \
|
||||
python examples/generate/multimodal/audio_language_offline.py \
|
||||
--model-type qwen2_5_omni
|
||||
|
||||
# Process image inputs
|
||||
python examples/offline_inference/vision_language.py \
|
||||
python examples/generate/multimodal/vision_language_offline.py \
|
||||
--modality image \
|
||||
--model-type qwen2_5_omni
|
||||
|
||||
# Process video inputs
|
||||
python examples/offline_inference/vision_language.py \
|
||||
python examples/generate/multimodal/vision_language_offline.py \
|
||||
--modality video \
|
||||
--model-type qwen2_5_omni
|
||||
```
|
||||
Executable → Regular
Executable → Regular
+1
-1
@@ -1402,7 +1402,7 @@ def run_mantis(questions: list[str], modality: str) -> ModelRequestData:
|
||||
# MiniCPM-V
|
||||
def run_minicpmv_base(questions: list[str], modality: str, model_name):
|
||||
assert modality in ["image", "video", "image+video"]
|
||||
# If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language.py` # noqa
|
||||
# If you want to use `MiniCPM-o-2_6` with audio inputs, check `audio_language_offline.py` # noqa
|
||||
|
||||
# 2.0
|
||||
# The official repo doesn't work yet, so we need to use a fork for now
|
||||
@@ -4,68 +4,74 @@
|
||||
import torch
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.config import PoolerConfig
|
||||
from vllm.inputs import TextPrompt
|
||||
from vllm.multimodal.utils import fetch_image
|
||||
|
||||
# Initialize model
|
||||
model = LLM(
|
||||
model="jinaai/jina-embeddings-v4-vllm-text-matching",
|
||||
runner="pooling",
|
||||
max_model_len=1024,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
|
||||
# Create text prompts
|
||||
text1 = "Ein wunderschöner Sonnenuntergang am Strand"
|
||||
text1_prompt = TextPrompt(prompt=f"Query: {text1}")
|
||||
def main():
|
||||
# Initialize model
|
||||
model = LLM(
|
||||
model="jinaai/jina-embeddings-v4-vllm-text-matching",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
runner="pooling",
|
||||
max_model_len=1024,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
|
||||
text2 = "浜辺に沈む美しい夕日"
|
||||
text2_prompt = TextPrompt(prompt=f"Query: {text2}")
|
||||
# Create text prompts
|
||||
text1 = "Ein wunderschöner Sonnenuntergang am Strand"
|
||||
text1_prompt = TextPrompt(prompt=f"Query: {text1}")
|
||||
|
||||
# Create image prompt
|
||||
image = fetch_image(
|
||||
"https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
|
||||
)
|
||||
image_prompt = TextPrompt(
|
||||
prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
|
||||
multi_modal_data={"image": image},
|
||||
)
|
||||
text2 = "浜辺に沈む美しい夕日"
|
||||
text2_prompt = TextPrompt(prompt=f"Query: {text2}")
|
||||
|
||||
# Encode all prompts
|
||||
prompts = [text1_prompt, text2_prompt, image_prompt]
|
||||
outputs = model.encode(prompts, pooling_task="token_embed")
|
||||
# Create image prompt
|
||||
image = fetch_image(
|
||||
"https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/eskimo.jpg" # noqa: E501
|
||||
)
|
||||
image_prompt = TextPrompt(
|
||||
prompt="<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|>\n", # noqa: E501
|
||||
multi_modal_data={"image": image},
|
||||
)
|
||||
|
||||
# Encode all prompts
|
||||
prompts = [text1_prompt, text2_prompt, image_prompt]
|
||||
outputs = model.encode(prompts, pooling_task="token_embed")
|
||||
|
||||
def get_embeddings(outputs):
|
||||
VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
|
||||
|
||||
embeddings = []
|
||||
for output in outputs:
|
||||
if VISION_START_TOKEN_ID in output.prompt_token_ids:
|
||||
# Gather only vision tokens
|
||||
img_start_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
|
||||
)[0][0]
|
||||
img_end_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
|
||||
)[0][0]
|
||||
embeddings_tensor = output.outputs.data.detach().clone()[
|
||||
img_start_pos : img_end_pos + 1
|
||||
]
|
||||
else:
|
||||
# Use all tokens for text-only prompts
|
||||
embeddings_tensor = output.outputs.data.detach().clone()
|
||||
|
||||
# Pool and normalize embeddings
|
||||
pooled_output = (
|
||||
embeddings_tensor.sum(dim=0, dtype=torch.float32)
|
||||
/ embeddings_tensor.shape[0]
|
||||
)
|
||||
embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
|
||||
return embeddings
|
||||
|
||||
embeddings = get_embeddings(outputs)
|
||||
|
||||
for embedding in embeddings:
|
||||
print(embedding.shape)
|
||||
|
||||
|
||||
def get_embeddings(outputs):
|
||||
VISION_START_TOKEN_ID, VISION_END_TOKEN_ID = 151652, 151653
|
||||
|
||||
embeddings = []
|
||||
for output in outputs:
|
||||
if VISION_START_TOKEN_ID in output.prompt_token_ids:
|
||||
# Gather only vision tokens
|
||||
img_start_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_START_TOKEN_ID
|
||||
)[0][0]
|
||||
img_end_pos = torch.where(
|
||||
torch.tensor(output.prompt_token_ids) == VISION_END_TOKEN_ID
|
||||
)[0][0]
|
||||
embeddings_tensor = output.outputs.data.detach().clone()[
|
||||
img_start_pos : img_end_pos + 1
|
||||
]
|
||||
else:
|
||||
# Use all tokens for text-only prompts
|
||||
embeddings_tensor = output.outputs.data.detach().clone()
|
||||
|
||||
# Pool and normalize embeddings
|
||||
pooled_output = (
|
||||
embeddings_tensor.sum(dim=0, dtype=torch.float32)
|
||||
/ embeddings_tensor.shape[0]
|
||||
)
|
||||
embeddings.append(torch.nn.functional.normalize(pooled_output, dim=-1))
|
||||
return embeddings
|
||||
|
||||
|
||||
embeddings = get_embeddings(outputs)
|
||||
|
||||
for embedding in embeddings:
|
||||
print(embedding.shape)
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
from argparse import Namespace
|
||||
|
||||
from vllm import LLM, EngineArgs
|
||||
from vllm.config import PoolerConfig
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
@@ -13,6 +14,7 @@ def parse_args():
|
||||
# Set example specific arguments
|
||||
parser.set_defaults(
|
||||
model="BAAI/bge-m3",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
runner="pooling",
|
||||
enforce_eager=True,
|
||||
)
|
||||
@@ -32,15 +34,6 @@ def main(args: Namespace):
|
||||
# You should pass runner="pooling" for embedding models
|
||||
llm = LLM(**vars(args))
|
||||
|
||||
# Generate embedding. The output is a list of EmbeddingRequestOutputs.
|
||||
outputs = llm.embed(prompts)
|
||||
|
||||
# Print the outputs.
|
||||
print("\nGenerated Outputs:\n" + "-" * 60)
|
||||
for prompt, output in zip(prompts, outputs):
|
||||
embeds = output.outputs.embedding
|
||||
print(len(embeds))
|
||||
|
||||
# Generate embedding for each token. The output is a list of PoolingRequestOutput.
|
||||
outputs = llm.encode(prompts, pooling_task="token_embed")
|
||||
|
||||
@@ -50,6 +43,20 @@ def main(args: Namespace):
|
||||
multi_vector = output.outputs.data
|
||||
print(multi_vector.shape)
|
||||
|
||||
query = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
# Generate scores.
|
||||
outputs = llm.score(query, documents)
|
||||
# Print the outputs.
|
||||
print("\nGenerated Outputs:\n" + "-" * 60)
|
||||
for document, output in zip(documents, outputs):
|
||||
score = output.outputs.score
|
||||
print(f"Pair: {[query, document]!r} \nScore: {score}")
|
||||
print("-" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
|
||||
@@ -7,10 +7,11 @@ Example online usage of Pooling API for multi vector retrieval.
|
||||
Run `vllm serve <model> --runner pooling`
|
||||
to start up the server in vLLM. e.g.
|
||||
|
||||
vllm serve BAAI/bge-m3
|
||||
vllm serve BAAI/bge-m3 --pooler-config.task token_embed
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import pprint
|
||||
|
||||
import requests
|
||||
import torch
|
||||
@@ -32,7 +33,8 @@ def parse_args():
|
||||
|
||||
|
||||
def main(args):
|
||||
api_url = f"http://{args.host}:{args.port}/pooling"
|
||||
pooling_url = f"http://{args.host}:{args.port}/pooling"
|
||||
score_url = f"http://{args.host}:{args.port}/score"
|
||||
model_name = args.model
|
||||
|
||||
prompts = [
|
||||
@@ -43,11 +45,23 @@ def main(args):
|
||||
]
|
||||
prompt = {"model": model_name, "input": prompts}
|
||||
|
||||
pooling_response = post_http_request(prompt=prompt, api_url=api_url)
|
||||
pooling_response = post_http_request(prompt=prompt, api_url=pooling_url)
|
||||
for output in pooling_response.json()["data"]:
|
||||
multi_vector = torch.tensor(output["data"])
|
||||
print(multi_vector.shape)
|
||||
|
||||
queries = "What is the capital of France?"
|
||||
documents = [
|
||||
"The capital of Brazil is Brasilia.",
|
||||
"The capital of France is Paris.",
|
||||
]
|
||||
prompt = {"model": model_name, "queries": queries, "documents": documents}
|
||||
score_response = post_http_request(prompt=prompt, api_url=score_url)
|
||||
print("\nPrompt when queries is string and documents is a list:")
|
||||
pprint.pprint(prompt)
|
||||
print("\nScore Response:")
|
||||
pprint.pprint(score_response.json())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
|
||||
@@ -34,10 +34,7 @@ def _run_vllm(vllm_runner):
|
||||
mode=CompilationMode.VLLM_COMPILE,
|
||||
cudagraph_mode=CUDAGraphMode.NONE,
|
||||
),
|
||||
# Phi-tiny-MoE uses SWA, whose admission cap is `cdiv(L, block_size) + 1`
|
||||
# at default block_size=16 — i.e. 17 blocks for max_model_len=256. Use
|
||||
# 32 for headroom.
|
||||
num_gpu_blocks_override=32,
|
||||
num_gpu_blocks_override=8,
|
||||
):
|
||||
pass
|
||||
|
||||
@@ -193,7 +190,7 @@ def _run_model(vllm_runner, spec: ModelStartupSpec):
|
||||
cudagraph_mode=CUDAGraphMode.NONE,
|
||||
pass_config=PassConfig(fuse_allreduce_rms=False),
|
||||
),
|
||||
num_gpu_blocks_override=16,
|
||||
num_gpu_blocks_override=8,
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
@@ -405,9 +405,6 @@ def test_should_split():
|
||||
(None, 0, 1, False, 2048, CUDAGraphMode.NONE, 0),
|
||||
# truncated to nearest multiple of 8 or 16
|
||||
(None, 257, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 256),
|
||||
# max_num_batched_tokens <= max_cudagraph_capture_size should always be
|
||||
# captured even if not landing on a 16-stride step
|
||||
(None, 2048, 1, False, 257, CUDAGraphMode.FULL_AND_PIECEWISE, 257),
|
||||
# max from list
|
||||
([1, 2, 4, 15], None, 1, False, 2048, CUDAGraphMode.FULL_AND_PIECEWISE, 15),
|
||||
# SP forces full-graph compilation, sizes are filtered by TP
|
||||
|
||||
@@ -41,3 +41,21 @@ def test_language_model_only_affects_model_hash():
|
||||
base_hash = ModelConfig(model).compute_hash()
|
||||
lm_only_hash = ModelConfig(model, language_model_only=True).compute_hash()
|
||||
assert base_hash != lm_only_hash
|
||||
|
||||
|
||||
def test_mm_encoder_fp8_scale_path_requires_fp8():
|
||||
with pytest.raises(ValueError, match="mm_encoder_attn_dtype"):
|
||||
MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json")
|
||||
|
||||
|
||||
def test_mm_encoder_attn_dtype_hash_updates(tmp_path):
|
||||
scale_file = tmp_path / "scales.json"
|
||||
scale_file.write_text("{}")
|
||||
base_hash = MultiModalConfig().compute_hash()
|
||||
fp8_hash = MultiModalConfig(mm_encoder_attn_dtype="fp8").compute_hash()
|
||||
fp8_static_hash = MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_path=str(scale_file),
|
||||
).compute_hash()
|
||||
assert base_hash != fp8_hash
|
||||
assert fp8_hash != fp8_static_hash
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for ``system_fingerprint`` construction."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.openai import fingerprint as fp
|
||||
|
||||
|
||||
def _cfg(tp=1, pp=1, dp=1, ep=False, digest="a3b21f94deadbeef"):
|
||||
c = SimpleNamespace(
|
||||
parallel_config=SimpleNamespace(
|
||||
tensor_parallel_size=tp,
|
||||
pipeline_parallel_size=pp,
|
||||
data_parallel_size=dp,
|
||||
enable_expert_parallel=ep,
|
||||
)
|
||||
)
|
||||
c.compute_hash = lambda: digest # type: ignore[attr-defined]
|
||||
return c
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset():
|
||||
fp.set_default_fingerprint_mode("full")
|
||||
yield
|
||||
fp.set_default_fingerprint_mode("full")
|
||||
|
||||
|
||||
def test_four_modes_produce_expected_shapes():
|
||||
from vllm import __version__ as v
|
||||
|
||||
cfg = _cfg(tp=8, ep=True)
|
||||
|
||||
assert fp.build_system_fingerprint(cfg, "full") == (f"vllm-{v}-tp8-ep-a3b21f94")
|
||||
assert fp.build_system_fingerprint(cfg, "hash") == f"vllm-{v}-a3b21f94"
|
||||
assert fp.build_system_fingerprint(cfg, "custom", "my-fp") == "my-fp"
|
||||
assert fp.build_system_fingerprint(cfg, "none") is None
|
||||
|
||||
|
||||
def test_full_mode_emits_only_non_trivial_parallelism():
|
||||
from vllm import __version__ as v
|
||||
|
||||
# Single-GPU: nothing between version and hash.
|
||||
assert fp.build_system_fingerprint(_cfg(), "full") == f"vllm-{v}-a3b21f94"
|
||||
# All parallelism axes.
|
||||
assert (
|
||||
fp.build_system_fingerprint(_cfg(tp=8, pp=2, dp=4, ep=True), "full")
|
||||
== f"vllm-{v}-tp8-pp2-dp4-ep-a3b21f94"
|
||||
)
|
||||
|
||||
|
||||
def test_get_respects_set_default():
|
||||
cfg = _cfg(tp=8)
|
||||
full = fp.get_system_fingerprint(cfg)
|
||||
assert full == fp.get_system_fingerprint(cfg)
|
||||
|
||||
fp.set_default_fingerprint_mode("hash")
|
||||
hashed = fp.get_system_fingerprint(cfg)
|
||||
assert hashed != full
|
||||
assert "tp8" not in hashed
|
||||
|
||||
fp.set_default_fingerprint_mode("custom", "deploy-42")
|
||||
assert fp.get_system_fingerprint(cfg) == "deploy-42"
|
||||
|
||||
fp.set_default_fingerprint_mode("none")
|
||||
assert fp.get_system_fingerprint(cfg) is None
|
||||
|
||||
|
||||
def test_compute_hash_failure_does_not_raise():
|
||||
cfg = _cfg()
|
||||
cfg.compute_hash = lambda: (_ for _ in ()).throw(RuntimeError("boom"))
|
||||
assert fp.build_system_fingerprint(cfg, "full").endswith("-nohash")
|
||||
assert fp.build_system_fingerprint(cfg, "hash").endswith("-nohash")
|
||||
@@ -1,13 +1,12 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.models.utils import softmax
|
||||
from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
|
||||
from vllm import LLM, ClassificationRequestOutput, PoolingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.tasks import PoolingTask
|
||||
|
||||
@@ -66,18 +65,6 @@ def test_list_prompts(llm: LLM):
|
||||
assert len(outputs[i].outputs.probs) == num_labels
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_token_classify(llm: LLM, caplog_vllm):
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
|
||||
assert len(outputs) == 1
|
||||
assert isinstance(outputs[0], PoolingRequestOutput)
|
||||
assert outputs[0].prompt_token_ids == prompt_token_ids
|
||||
assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_pooling_params(llm: LLM):
|
||||
def get_outputs(use_activation):
|
||||
@@ -110,10 +97,12 @@ def test_score_api(llm: LLM):
|
||||
llm.score("ping", "pong", use_tqdm=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask):
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_classify":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
else:
|
||||
err_msg = "Embedding API is not supported by this model.+"
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
|
||||
@@ -436,26 +436,7 @@ async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
|
||||
task = "token_classify"
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == 8
|
||||
assert len(poolings.data[0].data[0]) == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -469,8 +450,11 @@ async def test_pooling_not_supported(
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_classify":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
@@ -38,11 +37,11 @@ def llm():
|
||||
seed=0,
|
||||
attention_config=attention_config,
|
||||
)
|
||||
assert embedding_size == llm.model_config.embedding_size
|
||||
|
||||
yield weakref.proxy(llm)
|
||||
|
||||
del llm
|
||||
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@@ -74,16 +73,6 @@ def test_list_prompts(llm: LLM):
|
||||
assert len(outputs[i].outputs.embedding) == embedding_size
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_token_embed(llm: LLM, caplog_vllm):
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
outputs = llm.encode(prompt, pooling_task="token_embed", use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
|
||||
multi_vector = outputs[0].outputs.data
|
||||
assert multi_vector.shape == (11, 384)
|
||||
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_pooling_params(llm: LLM):
|
||||
def get_outputs(normalize):
|
||||
@@ -107,10 +96,14 @@ def test_pooling_params(llm: LLM):
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["token_classify", "classify", "plugin"])
|
||||
@pytest.mark.parametrize(
|
||||
"task", ["token_classify", "classify", "token_embed", "plugin"]
|
||||
)
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask):
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_embed":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
else:
|
||||
err_msg = "Classification API is not supported by this model.+"
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
|
||||
@@ -732,28 +732,9 @@ async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
task = "token_embed"
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": model_name,
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert len(poolings.data[0].data[0]) == 384
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
|
||||
@pytest.mark.parametrize(
|
||||
"task", ["classify", "token_classify", "token_embed", "plugin"]
|
||||
)
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -769,6 +750,8 @@ async def test_pooling_not_supported(
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_embed":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -452,25 +452,6 @@ async def test_pooling_classify(server: RemoteOpenAIServer):
|
||||
assert len(poolings.data[0].data) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pooling_token_classify(server: RemoteOpenAIServer):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"task": "token_classify",
|
||||
"input": input_text,
|
||||
"encoding_format": "float",
|
||||
},
|
||||
)
|
||||
|
||||
poolings = PoolingResponse.model_validate(response.json())
|
||||
|
||||
assert len(poolings.data) == 1
|
||||
assert len(poolings.data[0].data) == len(input_tokens)
|
||||
assert len(poolings.data[0].data[0]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rerank_max_tokens_per_doc(
|
||||
server: RemoteOpenAIServer,
|
||||
@@ -544,7 +525,7 @@ async def test_rerank_max_tokens_per_doc_validation(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "token_classify", "plugin"])
|
||||
async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
|
||||
response = requests.post(
|
||||
server.url_for("pooling"),
|
||||
@@ -558,6 +539,8 @@ async def test_pooling_not_supported(server: RemoteOpenAIServer, task: str):
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "token_classify":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
@@ -60,22 +59,19 @@ def test_token_ids_prompts(llm: LLM):
|
||||
|
||||
@pytest.mark.skip_global_cleanup
|
||||
def test_score_api(llm: LLM):
|
||||
err_msg = "Scoring API is only enabled for num_labels == 1."
|
||||
err_msg = "This model does not support the Scoring API."
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.score("ping", "pong", use_tqdm=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
|
||||
if task == "classify":
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "classify":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
else:
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
else:
|
||||
err_msg = "Embedding API is not supported by this model.+"
|
||||
err_msg = "Embedding API is not supported by this model.+"
|
||||
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
|
||||
@@ -50,7 +50,7 @@ async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: st
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["classify", "embed", "token_embed", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -63,9 +63,12 @@ async def test_pooling_not_supported(
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "classify":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import logging
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
@@ -64,15 +63,12 @@ def test_token_ids_prompts(llm: LLM):
|
||||
|
||||
@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
|
||||
def test_unsupported_tasks(llm: LLM, task: PoolingTask, caplog_vllm):
|
||||
if task == "embed":
|
||||
with caplog_vllm.at_level(level=logging.WARNING, logger="vllm"):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
assert "deprecated" in caplog_vllm.text
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "embed":
|
||||
err_msg = "Try switching the model's pooling_task via.+"
|
||||
else:
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
else:
|
||||
err_msg = "Classification API is not supported by this model.+"
|
||||
err_msg = "Classification API is not supported by this model.+"
|
||||
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
with pytest.raises(ValueError, match=err_msg):
|
||||
llm.encode(prompt, pooling_task=task, use_tqdm=False)
|
||||
|
||||
@@ -73,7 +73,7 @@ async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
|
||||
@pytest.mark.parametrize("task", ["embed", "classify", "token_classify", "plugin"])
|
||||
async def test_pooling_not_supported(
|
||||
server: RemoteOpenAIServer, model_name: str, task: str
|
||||
):
|
||||
@@ -86,9 +86,12 @@ async def test_pooling_not_supported(
|
||||
"task": task,
|
||||
},
|
||||
)
|
||||
assert response.json()["error"]["type"] == "BadRequestError"
|
||||
|
||||
if task == "plugin":
|
||||
err_msg = "No IOProcessor plugin installed."
|
||||
elif task == "embed":
|
||||
err_msg = "Try switching the model's pooling_task via"
|
||||
else:
|
||||
err_msg = f"Unsupported task: {task!r}"
|
||||
assert response.json()["error"]["message"].startswith(err_msg)
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the full FP8 ViT attention path (quantize -> cuDNN -> un-pad)."""
|
||||
|
||||
import contextlib
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
from vllm.utils.flashinfer import (
|
||||
is_flashinfer_cudnn_fp8_prefill_attn_supported,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
|
||||
def _has_flashinfer_cudnn() -> bool:
|
||||
"""Check if FlashInfer cuDNN backend is available."""
|
||||
try:
|
||||
from flashinfer.prefill import (
|
||||
cudnn_batch_prefill_with_kv_cache, # noqa: F401
|
||||
)
|
||||
|
||||
return True
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
|
||||
HEAD_DIMS = [72, 80]
|
||||
SEQ_LENS = [256]
|
||||
NUM_HEADS = [16]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _fp8_attention():
|
||||
"""Create FP8-enabled MMEncoderAttention via config."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN FP8 prefill attention not supported")
|
||||
|
||||
mm_config = MultiModalConfig(mm_encoder_attn_dtype="fp8")
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
# MMEncoderAttention reads torch.get_default_dtype() during init
|
||||
# to determine the output dtype. In real model loading this is bf16.
|
||||
old_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
yield
|
||||
|
||||
torch.set_default_dtype(old_dtype)
|
||||
|
||||
|
||||
def _build_cu_seqlens_and_meta(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
fp8_padded_hidden_size: int | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Build cu_seqlens, max_seqlen, sequence_lengths for a single sequence."""
|
||||
import numpy as np
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
|
||||
cu_seqlens_np = np.array([0, seq_len], dtype=np.int32)
|
||||
|
||||
sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_seqlens_np,
|
||||
torch.device("cuda"),
|
||||
)
|
||||
|
||||
max_seqlen = torch.tensor(
|
||||
MMEncoderAttention.compute_max_seqlen(
|
||||
AttentionBackendEnum.FLASHINFER, cu_seqlens_np
|
||||
),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_seqlens_np,
|
||||
num_heads * head_dim,
|
||||
1, # tp_size
|
||||
torch.device("cuda"),
|
||||
fp8_padded_hidden_size=fp8_padded_hidden_size,
|
||||
)
|
||||
|
||||
return cu_seqlens, max_seqlen, sequence_lengths
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (HAS_TRITON and _has_flashinfer_cudnn()),
|
||||
reason="Triton and FlashInfer cuDNN required",
|
||||
)
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
def test_fp8_attn_output_shape(
|
||||
head_dim: int,
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
_fp8_attention,
|
||||
) -> None:
|
||||
"""Verify FP8 attention produces correct output shape after un-padding."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
|
||||
attn = None
|
||||
with contextlib.suppress(ValueError, ImportError):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 MMEncoderAttention not available")
|
||||
assert attn is not None # mypy narrowing
|
||||
|
||||
# FP8 always needs fp8_padded_hidden_size for correct cu_seqlens
|
||||
fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
|
||||
|
||||
cu_seqlens, max_seqlen, sequence_lengths = _build_cu_seqlens_and_meta(
|
||||
seq_len, num_heads, head_dim, fp8_padded_hidden_size=fp8_padded_hidden_size
|
||||
)
|
||||
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
output = attn._forward_flashinfer(q, k, v, cu_seqlens, max_seqlen, sequence_lengths)
|
||||
|
||||
# Output should have original head_dim (un-padded)
|
||||
assert output.shape[-1] == head_dim
|
||||
assert output.dtype == torch.bfloat16
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (HAS_TRITON and _has_flashinfer_cudnn()),
|
||||
reason="Triton and FlashInfer cuDNN required",
|
||||
)
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
def test_fp8_vs_bf16_close(
|
||||
head_dim: int, seq_len: int, num_heads: int, _fp8_attention
|
||||
) -> None:
|
||||
"""FP8 attention output should be reasonably close to BF16 baseline."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
|
||||
torch.manual_seed(42)
|
||||
q = torch.randn(
|
||||
1,
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
# FP8 path
|
||||
attn_fp8 = None
|
||||
with contextlib.suppress(ValueError, ImportError):
|
||||
attn_fp8 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
if attn_fp8 is None or not attn_fp8.fp8_enabled:
|
||||
pytest.skip("FP8 MMEncoderAttention not available")
|
||||
assert attn_fp8 is not None # mypy narrowing
|
||||
|
||||
fp8_padded_hidden_size = num_heads * round_up(head_dim, 16)
|
||||
cu_seqlens, max_seqlen, seq_lengths = _build_cu_seqlens_and_meta(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
fp8_padded_hidden_size=fp8_padded_hidden_size,
|
||||
)
|
||||
|
||||
out_fp8 = attn_fp8._forward_flashinfer(
|
||||
q.clone(),
|
||||
k.clone(),
|
||||
v.clone(),
|
||||
cu_seqlens,
|
||||
max_seqlen,
|
||||
seq_lengths,
|
||||
)
|
||||
|
||||
# BF16 baseline (create non-FP8 attention by using scale=attn_fp8.scale
|
||||
# and calling the wrapper directly without FP8 quantization)
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
_get_flashinfer_workspace_buffer,
|
||||
)
|
||||
from vllm.v1.attention.ops.vit_attn_wrappers import (
|
||||
vit_flashinfer_wrapper,
|
||||
)
|
||||
|
||||
out_bf16 = vit_flashinfer_wrapper(
|
||||
q=q.clone(),
|
||||
k=k.clone(),
|
||||
v=v.clone(),
|
||||
scale=attn_fp8.scale,
|
||||
workspace_buffer=_get_flashinfer_workspace_buffer(),
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max_seqlen,
|
||||
sequence_lengths=seq_lengths,
|
||||
)
|
||||
|
||||
out_fp8_f = out_fp8.float()
|
||||
out_bf16_f = out_bf16.float()
|
||||
|
||||
abs_diff = (out_fp8_f - out_bf16_f).abs()
|
||||
abs_diff_flat = abs_diff.flatten()
|
||||
|
||||
# Relative diff (avoid division by zero)
|
||||
denom = out_bf16_f.abs().clamp(min=1e-6)
|
||||
rel_diff_flat = (abs_diff / denom).flatten()
|
||||
|
||||
cosine_sim = torch.nn.functional.cosine_similarity(
|
||||
out_fp8_f.flatten().unsqueeze(0),
|
||||
out_bf16_f.flatten().unsqueeze(0),
|
||||
).item()
|
||||
|
||||
pcts = [50, 90, 95, 99, 99.9]
|
||||
abs_pct = {p: torch.quantile(abs_diff_flat, p / 100).item() for p in pcts}
|
||||
rel_pct = {p: torch.quantile(rel_diff_flat, p / 100).item() for p in pcts}
|
||||
|
||||
print(f"\nFP8 vs BF16 (head_dim={head_dim}, seq_len={seq_len}):")
|
||||
print(f" cosine_sim={cosine_sim:.6f}")
|
||||
print(
|
||||
f" abs_diff: max={abs_diff_flat.max().item():.6f}, "
|
||||
f"mean={abs_diff_flat.mean().item():.6f}, "
|
||||
+ ", ".join(f"p{p}={abs_pct[p]:.6f}" for p in pcts)
|
||||
)
|
||||
print(
|
||||
f" rel_diff: max={rel_diff_flat.max().item():.6f}, "
|
||||
f"mean={rel_diff_flat.mean().item():.6f}, "
|
||||
+ ", ".join(f"p{p}={rel_pct[p]:.6f}" for p in pcts)
|
||||
)
|
||||
|
||||
assert abs_diff_flat.max().item() < 0.3, (
|
||||
f"FP8 vs BF16 max abs diff too large: {abs_diff_flat.max().item()}"
|
||||
)
|
||||
assert abs_diff_flat.mean().item() < 0.03, (
|
||||
f"FP8 vs BF16 mean abs diff too large: {abs_diff_flat.mean().item()}"
|
||||
)
|
||||
assert cosine_sim > 0.99, f"Cosine similarity too low: {cosine_sim:.6f}"
|
||||
@@ -0,0 +1,124 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the stride-aware FP8 quantization kernel with head_dim padding."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
|
||||
if HAS_TRITON:
|
||||
from vllm.kernels.triton.qkv_padded_fp8_quant import (
|
||||
quantize_fp8_pad_head_dim_triton,
|
||||
)
|
||||
|
||||
HEAD_DIMS = [72, 80, 128]
|
||||
SEQ_LENS = [64, 256]
|
||||
NUM_HEADS = [16]
|
||||
SCALES = [0.01, 0.1, 1.0]
|
||||
|
||||
|
||||
def _naive_fp8_quantize(
|
||||
tensor: torch.Tensor, scale: torch.Tensor, skip_scale: bool
|
||||
) -> torch.Tensor:
|
||||
"""Reference FP8 quantization in PyTorch."""
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
fp8_max = torch.finfo(fp8_dtype).max
|
||||
fp8_min = -fp8_max
|
||||
|
||||
x = tensor.float()
|
||||
if not skip_scale:
|
||||
x = x / scale.item()
|
||||
x = x.clamp(fp8_min, fp8_max)
|
||||
return x.to(fp8_dtype)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("head_dim", HEAD_DIMS)
|
||||
@pytest.mark.parametrize("seq_len", SEQ_LENS)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("scale_val", SCALES)
|
||||
def test_quantize_contiguous(
|
||||
head_dim: int, seq_len: int, num_heads: int, scale_val: float
|
||||
) -> None:
|
||||
"""Test quantization of contiguous 3D tensors."""
|
||||
torch.manual_seed(42)
|
||||
tensor = torch.randn(
|
||||
seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
scale = torch.tensor([scale_val], dtype=torch.float32, device="cuda").view(
|
||||
1, 1, 1, 1
|
||||
)
|
||||
|
||||
result = quantize_fp8_pad_head_dim_triton(tensor, scale)
|
||||
|
||||
padded_dim = (head_dim + 15) // 16 * 16
|
||||
assert result.shape == (seq_len, num_heads, padded_dim)
|
||||
assert result.is_contiguous()
|
||||
assert result.dtype == current_platform.fp8_dtype()
|
||||
|
||||
# Compare unpadded portion against reference
|
||||
ref = _naive_fp8_quantize(tensor, scale, skip_scale=False)
|
||||
torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
|
||||
|
||||
# Padded region should be zero
|
||||
if padded_dim > head_dim:
|
||||
assert (result[:, :, head_dim:].float() == 0).all()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
@pytest.mark.parametrize("head_dim", [72, 80])
|
||||
def test_quantize_non_contiguous(head_dim: int) -> None:
|
||||
"""Test quantization from non-contiguous QKV views (interleaved buffer)."""
|
||||
seq_len, num_heads = 64, 16
|
||||
# Simulate interleaved QKV buffer: shape (seq_len, 3 * num_heads, head_dim)
|
||||
qkv = torch.randn(
|
||||
seq_len, 3 * num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
# Q is every 3rd head slice - non-contiguous view
|
||||
q = qkv[:, 0::3, :]
|
||||
assert not q.is_contiguous()
|
||||
|
||||
scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
result = quantize_fp8_pad_head_dim_triton(q, scale)
|
||||
|
||||
padded_dim = (head_dim + 15) // 16 * 16
|
||||
assert result.shape == (seq_len, num_heads, padded_dim)
|
||||
assert result.is_contiguous()
|
||||
|
||||
# Compare against contiguous reference
|
||||
ref = _naive_fp8_quantize(q.contiguous(), scale, skip_scale=False)
|
||||
torch.testing.assert_close(result[:, :, :head_dim].float(), ref.float())
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
def test_skip_scale() -> None:
|
||||
"""Test skip_scale=True produces cast-only output (no division)."""
|
||||
seq_len, num_heads, head_dim = 32, 8, 80
|
||||
tensor = torch.randn(
|
||||
seq_len, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
scale = torch.tensor([0.5], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
|
||||
result_skip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=True)
|
||||
result_noskip = quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=False)
|
||||
|
||||
# skip_scale should just cast, not divide
|
||||
ref_cast = _naive_fp8_quantize(tensor, scale, skip_scale=True)
|
||||
torch.testing.assert_close(result_skip[:, :, :head_dim].float(), ref_cast.float())
|
||||
|
||||
# With scale != 1.0, skip and no-skip should differ
|
||||
assert not torch.equal(result_skip.float(), result_noskip.float())
|
||||
|
||||
|
||||
@pytest.mark.skipif(not HAS_TRITON, reason="Triton not available")
|
||||
def test_4d_input() -> None:
|
||||
"""Test that 4D input (B, S, H, D) is handled correctly."""
|
||||
B, S, H, D = 2, 32, 8, 72
|
||||
tensor = torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16)
|
||||
scale = torch.tensor([0.1], dtype=torch.float32, device="cuda").view(1, 1, 1, 1)
|
||||
|
||||
result = quantize_fp8_pad_head_dim_triton(tensor, scale)
|
||||
padded_dim = (D + 15) // 16 * 16
|
||||
assert result.shape == (B, S, H, padded_dim)
|
||||
@@ -0,0 +1,251 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for FP8 scaling (dynamic and static) in MMEncoderAttention."""
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
_FP8_AMAX_HISTORY_LEN,
|
||||
_FP8_MAX,
|
||||
)
|
||||
from vllm.utils.flashinfer import (
|
||||
is_flashinfer_cudnn_fp8_prefill_attn_supported,
|
||||
)
|
||||
|
||||
LAYER_0 = "visual.blocks.0.attn.attn"
|
||||
LAYER_1 = "visual.blocks.1.attn.attn"
|
||||
NUM_HEADS = 16
|
||||
HEAD_DIM = 72
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _build_attention(mm_config):
|
||||
"""Yield an MMEncoderAttention with the given multimodal config.
|
||||
|
||||
The VllmConfig context stays active while the test runs so that
|
||||
``get_multimodal_config()`` calls during the forward path resolve. Also
|
||||
invokes ``process_weights_after_loading`` to simulate the model loader's
|
||||
auto-scan. Yields ``None`` if FlashInfer cuDNN is not available.
|
||||
"""
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
yield None
|
||||
return
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=NUM_HEADS,
|
||||
head_size=HEAD_DIM,
|
||||
prefix=LAYER_0,
|
||||
)
|
||||
attn.process_weights_after_loading(torch.bfloat16)
|
||||
yield attn
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _make_attention():
|
||||
"""Create an MMEncoderAttention with dynamic FP8 scaling."""
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
with _build_attention(MultiModalConfig(mm_encoder_attn_dtype="fp8")) as attn:
|
||||
yield attn
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _make_static_attention(tmp_path):
|
||||
"""Create an MMEncoderAttention with static FP8 scales from a file."""
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
scale_file = tmp_path / "scales.json"
|
||||
scale_file.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
LAYER_0: {"q": 224.0, "k": 198.0, "v": 210.0},
|
||||
LAYER_1: {"q": 100.0, "k": 110.0, "v": 120.0},
|
||||
}
|
||||
)
|
||||
)
|
||||
with _build_attention(
|
||||
MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_path=str(scale_file),
|
||||
)
|
||||
) as attn:
|
||||
yield attn
|
||||
|
||||
|
||||
def test_dynamic_scaling_updates_scales(_make_attention) -> None:
|
||||
"""Verify that _record_amax_and_update_scales updates scale buffers."""
|
||||
attn = _make_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
|
||||
S, H, D = 32, NUM_HEADS, HEAD_DIM
|
||||
q = torch.full((S, H, D), 2.0, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), 3.0, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), 4.0, device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
expected_q_scale = 2.0 / _FP8_MAX
|
||||
expected_k_scale = 3.0 / _FP8_MAX
|
||||
expected_v_scale = 4.0 / _FP8_MAX
|
||||
|
||||
torch.testing.assert_close(attn._fp8_q_scale.item(), expected_q_scale)
|
||||
torch.testing.assert_close(attn._fp8_k_scale.item(), expected_k_scale)
|
||||
torch.testing.assert_close(attn._fp8_v_scale.item(), expected_v_scale)
|
||||
|
||||
|
||||
def test_circular_buffer_wraps(_make_attention) -> None:
|
||||
"""Verify the amax circular buffer wraps at HISTORY_LEN."""
|
||||
attn = _make_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
S, H, D = 16, NUM_HEADS, HEAD_DIM
|
||||
|
||||
for i in range(_FP8_AMAX_HISTORY_LEN + 2):
|
||||
mag = float(i + 1)
|
||||
q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
assert attn._fp8_amax_pos == 2
|
||||
|
||||
expected_max = float(_FP8_AMAX_HISTORY_LEN + 2)
|
||||
expected_scale = expected_max / _FP8_MAX
|
||||
torch.testing.assert_close(attn._fp8_q_scale.item(), expected_scale)
|
||||
|
||||
|
||||
def test_static_scales_loaded(_make_static_attention) -> None:
|
||||
"""Verify static scales are loaded from the JSON file."""
|
||||
attn = _make_static_attention
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available (FlashInfer backend required)")
|
||||
|
||||
assert attn.fp8_enabled
|
||||
assert not attn._fp8_dynamic_scale
|
||||
|
||||
# Layer 0 scales (the layer this attention was created with).
|
||||
assert attn._fp8_q_scale.item() == 224.0
|
||||
assert attn._fp8_k_scale.item() == 198.0
|
||||
assert attn._fp8_v_scale.item() == 210.0
|
||||
|
||||
assert not attn.skip_scale_q
|
||||
assert not attn.skip_scale_k
|
||||
assert not attn.skip_scale_v
|
||||
|
||||
# No amax history buffers for static scaling.
|
||||
assert not hasattr(attn, "_fp8_q_amax")
|
||||
|
||||
|
||||
def test_static_scales_missing_layer(tmp_path) -> None:
|
||||
"""Verify error when requested layer is not in the scale file."""
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN not available")
|
||||
|
||||
scale_file = tmp_path / "wrong_layer.json"
|
||||
scale_file.write_text(
|
||||
json.dumps({"visual.blocks.99.attn": {"q": 1.0, "k": 1.0, "v": 1.0}})
|
||||
)
|
||||
mm_config = MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_path=str(scale_file),
|
||||
)
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.model_config = SimpleNamespace(multimodal_config=mm_config)
|
||||
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
),
|
||||
):
|
||||
attn = MMEncoderAttention(
|
||||
num_heads=NUM_HEADS,
|
||||
head_size=HEAD_DIM,
|
||||
prefix=LAYER_0,
|
||||
)
|
||||
with pytest.raises(ValueError, match="scales not found for layer"):
|
||||
attn.process_weights_after_loading(torch.bfloat16)
|
||||
|
||||
|
||||
def test_dynamic_scales_auto_save(tmp_path) -> None:
|
||||
"""Verify scales are saved to disk after the amax buffer fills."""
|
||||
import vllm.model_executor.layers.attention.mm_encoder_attention as _mod
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
|
||||
if not is_flashinfer_cudnn_fp8_prefill_attn_supported():
|
||||
pytest.skip("FlashInfer cuDNN not available")
|
||||
|
||||
# Reset module-level state between runs (other tests may have left
|
||||
# state behind after triggering a save).
|
||||
_mod._fp8_scale_save_path = None
|
||||
_mod._fp8_saved_scale_refs.clear()
|
||||
|
||||
save_file = tmp_path / "auto_scales.json"
|
||||
with _build_attention(
|
||||
MultiModalConfig(
|
||||
mm_encoder_attn_dtype="fp8",
|
||||
mm_encoder_fp8_scale_save_path=str(save_file),
|
||||
)
|
||||
) as attn:
|
||||
if attn is None or not attn.fp8_enabled:
|
||||
pytest.skip("FP8 attention not available")
|
||||
|
||||
attn = attn.to("cuda")
|
||||
S, H, D = 16, NUM_HEADS, HEAD_DIM
|
||||
|
||||
# Run exactly _FP8_AMAX_HISTORY_LEN forward passes.
|
||||
for i in range(_FP8_AMAX_HISTORY_LEN):
|
||||
mag = float(i + 1)
|
||||
q = torch.full((S, H, D), mag, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.full((S, H, D), mag * 0.5, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.full((S, H, D), mag * 0.3, device="cuda", dtype=torch.bfloat16)
|
||||
attn._record_amax_and_update_scales(q, k, v)
|
||||
|
||||
# File should have been written on the 16th call (buffer wrap).
|
||||
assert save_file.is_file(), "Scale file was not saved"
|
||||
scales = json.loads(save_file.read_text())
|
||||
assert LAYER_0 in scales
|
||||
assert set(scales[LAYER_0].keys()) == {"q", "k", "v"}
|
||||
for val in scales[LAYER_0].values():
|
||||
assert isinstance(val, float) and val > 0
|
||||
|
||||
# Path is cleared after the one-shot save fires.
|
||||
assert _mod._fp8_scale_save_path is None
|
||||
@@ -3,11 +3,12 @@
|
||||
"""
|
||||
Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.
|
||||
|
||||
Four test functions cover five paths:
|
||||
Two paths tested:
|
||||
A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
|
||||
B) Indexer: head_dim=128 (all FP8), quant_block=128
|
||||
C) DeepseekV4 Attention magnitude range: correctness across small/large values
|
||||
D) Indexer fused Triton kernel: compress+norm+rope+quant+insert
|
||||
|
||||
These serve as golden references for validating the future fused
|
||||
compressor+quant+cache kernel.
|
||||
"""
|
||||
|
||||
import math
|
||||
@@ -20,12 +21,6 @@ from vllm.v1.attention.ops.deepseek_v4_ops import (
|
||||
dequantize_and_gather_k_cache,
|
||||
quantize_and_insert_k_cache,
|
||||
)
|
||||
from vllm.v1.attention.ops.deepseek_v4_ops.fused_compress_quant_cache import (
|
||||
_fused_kv_compress_norm_rope_insert_indexer_attn,
|
||||
_fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
|
||||
)
|
||||
|
||||
from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
|
||||
|
||||
|
||||
def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
|
||||
@@ -314,222 +309,3 @@ def test_deepseek_v4_quant_magnitude_range():
|
||||
f"Token {t}: rel_err={rel_err:.4f}, abs_diff={abs_diff:.6f}, "
|
||||
f"magnitude={magnitude:.4f}"
|
||||
)
|
||||
|
||||
|
||||
# ── Test D: Indexer fused K-cache insert (Triton kernels) ────────────────────
|
||||
#
|
||||
# Both kernels share the same Triton signature; use_fp4 selects between them.
|
||||
# Full pipeline: state-cache gather → softmax-weighted compress → RMSNorm →
|
||||
# GPT-J RoPE → quant (MXFP4 or FP8) → paged cache insert.
|
||||
|
||||
|
||||
def _reference_kv_compress_norm_rope(
|
||||
state_cache: torch.Tensor,
|
||||
block_table: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
rms_weight: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
compress_ratio: int = 1,
|
||||
overlap: int = 0,
|
||||
use_fp4: bool = False,
|
||||
rms_eps: float = 1e-6,
|
||||
fp8_max: float = 448.0,
|
||||
):
|
||||
"""Compress → RMSNorm → GPT-J RoPE → quantize.
|
||||
|
||||
Gathers (1+overlap)*compress_ratio state entries per output token, applies
|
||||
per-element softmax over the scores, and computes the weighted kv sum.
|
||||
Returns (quantized_values, scale) matching the kernel's output layout.
|
||||
"""
|
||||
device = state_cache.device
|
||||
head_dim = rms_weight.shape[0]
|
||||
rope_dim = cos_sin_cache.shape[-1]
|
||||
state_block_size = state_cache.shape[1]
|
||||
state_width = state_cache.shape[-1] // 2
|
||||
nope_dim = head_dim - rope_dim
|
||||
total = (1 + overlap) * compress_ratio
|
||||
results = []
|
||||
for pos in positions.tolist():
|
||||
src = torch.arange(pos - total + 1, pos + 1, dtype=torch.int64, device=device)
|
||||
valid = src >= 0
|
||||
idx = src.clamp(min=0)
|
||||
pages = block_table[0, idx // state_block_size]
|
||||
offsets = idx % state_block_size
|
||||
raw = state_cache[pages, offsets].float() # [total, state_dim]
|
||||
|
||||
# Group 0 (tokens 0..cr-1): kv[:H], score[SW:SW+H]
|
||||
# Group 1 (tokens cr..2cr-1): kv[H:2H], score[SW+H:SW+2H]
|
||||
if overlap:
|
||||
sw = state_width
|
||||
g0_kv = raw[:compress_ratio, :head_dim]
|
||||
g1_kv = raw[compress_ratio:, head_dim : 2 * head_dim]
|
||||
g0_scores = raw[:compress_ratio, sw : sw + head_dim]
|
||||
g1_scores = raw[compress_ratio:, sw + head_dim : sw + 2 * head_dim]
|
||||
kv = torch.cat([g0_kv, g1_kv])
|
||||
scores = torch.cat([g0_scores, g1_scores])
|
||||
else:
|
||||
kv = raw[:, :head_dim]
|
||||
scores = raw[:, state_width : state_width + head_dim]
|
||||
|
||||
scores[~valid] = float("-inf")
|
||||
kv[~valid] = 0.0
|
||||
weights = torch.softmax(scores, dim=0)
|
||||
compressed = (kv * weights).sum(dim=0) # [H]
|
||||
var = (compressed * compressed).mean()
|
||||
normed = compressed * torch.rsqrt(var + rms_eps) * rms_weight.float()
|
||||
compressed_pos = (pos // compress_ratio) * compress_ratio
|
||||
cos, sin = cos_sin_cache[compressed_pos].float().chunk(2)
|
||||
nope, rope = normed.split([nope_dim, rope_dim])
|
||||
rope = torch.stack(
|
||||
[rope[0::2] * cos - rope[1::2] * sin, rope[1::2] * cos + rope[0::2] * sin],
|
||||
dim=-1,
|
||||
).reshape(rope_dim)
|
||||
results.append(torch.cat([nope, rope]).to(state_cache.dtype))
|
||||
result = torch.stack(results)
|
||||
|
||||
if use_fp4:
|
||||
return quantize_to_mxfp4(result)
|
||||
else:
|
||||
pairs = [
|
||||
_ue8m0_reference(result[t], head_dim, fp8_max) for t in range(len(result))
|
||||
]
|
||||
quants, scales = zip(*pairs)
|
||||
return torch.stack(quants), torch.cat(scales)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 7, 32])
|
||||
@pytest.mark.parametrize("kv_block_size", [16, 32])
|
||||
@pytest.mark.parametrize("use_fp4", [False, True])
|
||||
def test_fused_kv_insert_indexer(num_tokens: int, kv_block_size: int, use_fp4: bool):
|
||||
"""Fused K compress+norm+rope+quant+insert for the indexer KV cache."""
|
||||
HEAD_DIM = 128
|
||||
ROPE_DIM = 64
|
||||
BLOCK_SIZE = 16
|
||||
RMS_EPS = 1e-6
|
||||
FP8_MAX = 448.0
|
||||
|
||||
device = "cuda"
|
||||
torch.manual_seed(42)
|
||||
compress_ratio = 4
|
||||
|
||||
if use_fp4:
|
||||
TOKEN_STRIDE = HEAD_DIM // 2 # packed nibbles: 64 bytes
|
||||
SCALE_DIM = HEAD_DIM // 32 # ue8m0 bytes: 4
|
||||
QUANT_BLOCK = 32
|
||||
kernel = _fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn
|
||||
else:
|
||||
TOKEN_STRIDE = HEAD_DIM # FP8 bytes: 128
|
||||
SCALE_DIM = 4 # 1 float32: 4 bytes
|
||||
QUANT_BLOCK = HEAD_DIM
|
||||
kernel = _fused_kv_compress_norm_rope_insert_indexer_attn
|
||||
|
||||
# overlap=1 whenever compress_ratio==4, matching DeepseekCompressor logic.
|
||||
overlap = 1 if compress_ratio == 4 else 0
|
||||
coff = 1 + overlap # multiplier for state_dim per entry
|
||||
|
||||
num_pages = (compress_ratio * num_tokens - 1) // BLOCK_SIZE + 2
|
||||
state_cache = torch.randn(
|
||||
num_pages,
|
||||
BLOCK_SIZE,
|
||||
2 * coff * HEAD_DIM, # kv_state + score_state, each coff*HEAD_DIM wide
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
block_table = torch.arange(num_pages, dtype=torch.int32, device=device).unsqueeze(0)
|
||||
token_to_req = torch.zeros(num_tokens, dtype=torch.int32, device=device)
|
||||
slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device)
|
||||
positions = torch.arange(
|
||||
compress_ratio - 1,
|
||||
compress_ratio * num_tokens,
|
||||
compress_ratio,
|
||||
dtype=torch.int64,
|
||||
device=device,
|
||||
)
|
||||
rms_weight = torch.randn(HEAD_DIM, dtype=torch.bfloat16, device=device)
|
||||
cos_sin_cache = torch.randn(compress_ratio * num_tokens, ROPE_DIM, device=device)
|
||||
|
||||
kv_n_blocks = (num_tokens + kv_block_size - 1) // kv_block_size + 1
|
||||
kv_cache = torch.zeros(
|
||||
kv_n_blocks,
|
||||
kv_block_size * (TOKEN_STRIDE + SCALE_DIM),
|
||||
dtype=torch.uint8,
|
||||
device=device,
|
||||
)
|
||||
|
||||
kernel[(num_tokens,)](
|
||||
state_cache,
|
||||
state_cache.stride(0),
|
||||
state_cache.stride(1),
|
||||
token_to_req,
|
||||
positions,
|
||||
slot_mapping,
|
||||
block_table,
|
||||
block_table.stride(0),
|
||||
BLOCK_SIZE,
|
||||
rms_weight,
|
||||
RMS_EPS,
|
||||
cos_sin_cache,
|
||||
cos_sin_cache.stride(0),
|
||||
kv_cache,
|
||||
slot_mapping,
|
||||
kv_block_size,
|
||||
HEAD_SIZE=HEAD_DIM,
|
||||
TRITON_BLOCK_SIZE=HEAD_DIM,
|
||||
STATE_WIDTH=coff * HEAD_DIM,
|
||||
COMPRESS_RATIO=compress_ratio,
|
||||
OVERLAP=overlap,
|
||||
ROPE_HEAD_DIM=ROPE_DIM,
|
||||
FP8_MAX=FP8_MAX,
|
||||
QUANT_BLOCK=QUANT_BLOCK,
|
||||
TOKEN_STRIDE=TOKEN_STRIDE,
|
||||
SCALE_DIM=SCALE_DIM,
|
||||
KV_BLOCK_STRIDE=kv_cache.stride(0),
|
||||
num_warps=1,
|
||||
)
|
||||
|
||||
k_quant, scale = _reference_kv_compress_norm_rope(
|
||||
state_cache,
|
||||
block_table,
|
||||
positions,
|
||||
rms_weight,
|
||||
cos_sin_cache,
|
||||
compress_ratio,
|
||||
overlap,
|
||||
use_fp4,
|
||||
rms_eps=RMS_EPS,
|
||||
fp8_max=FP8_MAX,
|
||||
)
|
||||
|
||||
if use_fp4:
|
||||
for i in range(num_tokens):
|
||||
blk, pos = i // kv_block_size, i % kv_block_size
|
||||
val_off = pos * TOKEN_STRIDE
|
||||
fp4_actual = kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
|
||||
assert torch.equal(k_quant[i], fp4_actual), (
|
||||
f"token {i}: packed nibbles differ, "
|
||||
f"{(k_quant[i] != fp4_actual).sum()} "
|
||||
f"/ {TOKEN_STRIDE}"
|
||||
)
|
||||
|
||||
scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
|
||||
scale_actual = kv_cache[blk, scale_off : scale_off + SCALE_DIM]
|
||||
assert torch.equal(scale_actual, scale[i]), (
|
||||
f"token {i}: ue8m0 {scale_actual.tolist()} != {scale[i].tolist()}"
|
||||
)
|
||||
|
||||
else:
|
||||
k_quant = k_quant.view(torch.uint8)
|
||||
for i in range(num_tokens):
|
||||
blk, pos = i // kv_block_size, i % kv_block_size
|
||||
val_off = pos * TOKEN_STRIDE
|
||||
assert torch.equal(
|
||||
k_quant[i], kv_cache[blk, val_off : val_off + TOKEN_STRIDE]
|
||||
), f"token {i}: FP8 bytes differ"
|
||||
|
||||
scale_off = kv_block_size * TOKEN_STRIDE + pos * SCALE_DIM
|
||||
actual_scale = kv_cache[blk, scale_off : scale_off + SCALE_DIM].view(
|
||||
torch.float32
|
||||
)
|
||||
assert torch.equal(actual_scale, scale[i : i + 1]), (
|
||||
f"token {i}: scale {actual_scale.item()} != {scale[i].item()}"
|
||||
)
|
||||
|
||||
@@ -30,56 +30,6 @@ N_HEAD = 64
|
||||
MAX_POS = 4096
|
||||
|
||||
|
||||
def quantize_to_mxfp4(
|
||||
x: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Reference MXFP4 quantization.
|
||||
|
||||
Args:
|
||||
x: [..., head_dim] where head_dim is divisible by 32
|
||||
Returns:
|
||||
packed: [..., head_dim//2] uint8 2 E2M1 nibbles/byte, low nibble = even index
|
||||
scales: [..., head_dim//32] uint8 1 ue8m0 byte
|
||||
"""
|
||||
MXFP4_BLOCK_SIZE = 32
|
||||
orig_shape = x.shape
|
||||
head_dim = orig_shape[-1]
|
||||
n_blocks = head_dim // MXFP4_BLOCK_SIZE
|
||||
|
||||
x_f32 = x.float().reshape(-1, n_blocks, MXFP4_BLOCK_SIZE)
|
||||
|
||||
# Per-block ue8m0 scale: 2^ceil(log2(amax / 6.0)), stored as byte = exp + 127
|
||||
# 6 * 2^-126 is from https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/kernel.py#L163
|
||||
amax = x_f32.abs().amax(dim=-1, keepdim=True).clamp(min=6 * (2**-126))
|
||||
log2_ratio = (amax * (1.0 / 6.0)).log2().ceil().clamp(-127.0, 127.0)
|
||||
scale = log2_ratio.exp2()
|
||||
ue8m0 = (log2_ratio + 127.0).to(torch.uint8) # [*, n_blocks]
|
||||
|
||||
# E2M1 round-to-nearest-even: midpoints round to the even code.
|
||||
# E2M1 values: [0.00, 0.50, 1.00, 1.50, 2.00, 3.00, 4.00, 6.00]
|
||||
# boundaries: [ 0.25, 0.75, 1.25, 1.75, 2.50, 3.50, 5.00]
|
||||
x_scaled = (x_f32 / scale).clamp(-6.0, 6.0)
|
||||
abs_x = x_scaled.abs()
|
||||
code = torch.zeros_like(abs_x, dtype=torch.int32)
|
||||
code = torch.where(abs_x > 0.25, 1, code)
|
||||
code = torch.where(abs_x >= 0.75, 2, code)
|
||||
code = torch.where(abs_x > 1.25, 3, code)
|
||||
code = torch.where(abs_x >= 1.75, 4, code)
|
||||
code = torch.where(abs_x > 2.5, 5, code)
|
||||
code = torch.where(abs_x >= 3.5, 6, code)
|
||||
code = torch.where(abs_x > 5.0, 7, code)
|
||||
sign = ((x_scaled.view(torch.int32) >> 31) & 1).to(torch.uint8)
|
||||
nibble = code.to(torch.uint8) | (sign << 3)
|
||||
|
||||
# Pack: even-index element → low nibble, odd-index → high nibble
|
||||
nibble_flat = nibble.reshape(-1, head_dim)
|
||||
packed = (nibble_flat[:, 0::2] | (nibble_flat[:, 1::2] << 4)).contiguous()
|
||||
packed = packed.reshape(*orig_shape[:-1], head_dim // 2)
|
||||
|
||||
scales = ue8m0.view(*orig_shape[:-1], n_blocks)
|
||||
return packed, scales
|
||||
|
||||
|
||||
def _reference(
|
||||
positions: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
@@ -87,7 +37,6 @@ def _reference(
|
||||
weights: torch.Tensor,
|
||||
softmax_scale: float,
|
||||
head_scale: float,
|
||||
use_fp4: bool = False,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
q_rot = q.clone()
|
||||
ops.rotary_embedding(
|
||||
@@ -100,33 +49,22 @@ def _reference(
|
||||
HEAD_DIM - ROPE_DIM, # rope_dim_offset → rotate the tail
|
||||
False,
|
||||
)
|
||||
q_fp8, q_scale = per_token_group_quant_fp8(
|
||||
q_rot.view(-1, HEAD_DIM).contiguous(),
|
||||
HEAD_DIM,
|
||||
use_ue8m0=True,
|
||||
)
|
||||
q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
|
||||
q_scale = q_scale.view(-1, N_HEAD)
|
||||
|
||||
if use_fp4:
|
||||
q_packed, ue8m0 = quantize_to_mxfp4(q_rot.view(-1, N_HEAD, HEAD_DIM))
|
||||
# Pack 4 ue8m0 bytes into 1 int32
|
||||
q_scale = ue8m0.view(torch.int32).squeeze(-1)
|
||||
# FP4 path: q_scale stays separate (cannot be folded into a per-token scalar)
|
||||
weights_out = weights.to(torch.float32) * softmax_scale * head_scale
|
||||
return (q_packed, q_scale), weights_out
|
||||
|
||||
else:
|
||||
q_fp8, q_scale = per_token_group_quant_fp8(
|
||||
q_rot.view(-1, HEAD_DIM).contiguous(),
|
||||
HEAD_DIM,
|
||||
use_ue8m0=True,
|
||||
)
|
||||
q_fp8 = q_fp8.view(-1, N_HEAD, HEAD_DIM)
|
||||
q_scale = q_scale.view(-1, N_HEAD)
|
||||
|
||||
weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
|
||||
return q_fp8, weights_out
|
||||
weights_out = weights.to(torch.float32) * q_scale * softmax_scale * head_scale
|
||||
return q_fp8, weights_out
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 7, 32, 257])
|
||||
@pytest.mark.parametrize("cache_dtype", [torch.float32, torch.bfloat16])
|
||||
@pytest.mark.parametrize("use_fp4", [False, True])
|
||||
@torch.inference_mode()
|
||||
def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use_fp4):
|
||||
def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype):
|
||||
device = "cuda"
|
||||
torch.manual_seed(0)
|
||||
|
||||
@@ -139,32 +77,21 @@ def test_fused_indexer_q_rope_quant_matches_unfused(num_tokens, cache_dtype, use
|
||||
softmax_scale = HEAD_DIM**-0.5
|
||||
head_scale = N_HEAD**-0.5
|
||||
|
||||
q_quant_ref, weights_ref = _reference(
|
||||
positions, q, cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
|
||||
q_fp8_ref, weights_ref = _reference(
|
||||
positions, q, cos_sin_cache, weights, softmax_scale, head_scale
|
||||
)
|
||||
q_quant_fused, weights_fused = fused_indexer_q_rope_quant(
|
||||
positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale, use_fp4
|
||||
q_fp8_fused, weights_fused = fused_indexer_q_rope_quant(
|
||||
positions, q.clone(), cos_sin_cache, weights, softmax_scale, head_scale
|
||||
)
|
||||
|
||||
if use_fp4:
|
||||
q_quant_ref, q_scale_ref = q_quant_ref
|
||||
q_quant_fused, q_scale_fused = q_quant_fused
|
||||
|
||||
assert torch.equal(q_scale_ref, q_scale_fused), (
|
||||
f"q_scale mismatch: "
|
||||
f"{(q_scale_ref != q_scale_fused).sum().item()} "
|
||||
f"/ {q_scale_ref.numel()} bytes differ"
|
||||
)
|
||||
|
||||
# fp8 tensors aren't directly comparable via torch.equal — reinterpret as int8.
|
||||
ref_bits = q_quant_ref.view(torch.int8)
|
||||
fused_bits = q_quant_fused.view(torch.int8)
|
||||
ref_bits = q_fp8_ref.view(torch.int8)
|
||||
fused_bits = q_fp8_fused.view(torch.int8)
|
||||
assert torch.equal(ref_bits, fused_bits), (
|
||||
f"q_quant_fused mismatch: "
|
||||
f"q_fp8 mismatch: "
|
||||
f"{(ref_bits != fused_bits).sum().item()} / {ref_bits.numel()} bytes differ"
|
||||
)
|
||||
|
||||
assert weights_fused.dtype == torch.float32
|
||||
assert torch.equal(weights_ref, weights_fused), (
|
||||
f"weights mismatch: max abs diff "
|
||||
f"{(weights_ref - weights_fused).abs().max().item()}"
|
||||
|
||||
@@ -6,6 +6,7 @@ from transformers import AutoModel
|
||||
|
||||
from tests.models.utils import check_embeddings_close
|
||||
from vllm import TokensPrompt
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -21,6 +22,7 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
max_model_len=128,
|
||||
max_num_batched_tokens=chunk_size,
|
||||
enforce_eager=True,
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import torch
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
@@ -25,29 +24,42 @@ sentences_2 = [
|
||||
similarity_reference = [[0.6259, 0.3474], [0.3309, 0.6734]]
|
||||
lexical_score_reference = [0.19554901123046875, 0.0]
|
||||
colbert_score_reference = [0.7797, 0.4620]
|
||||
SUPPORTED_TASKS = ["embed", "token_embed", "token_classify"]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=SUPPORTED_TASKS)
|
||||
def pooling_task(request):
|
||||
yield request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
def server(pooling_task):
|
||||
args = [
|
||||
"--max-model-len",
|
||||
str(MAX_MODEL_LEN),
|
||||
"--hf-overrides",
|
||||
'{"architectures": ["BgeM3EmbeddingModel"]}',
|
||||
"--pooler-config.task",
|
||||
pooling_task,
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server):
|
||||
async with server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_embedding(client: openai.AsyncOpenAI):
|
||||
async def test_bge_m3_api_server_embedding(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "embed":
|
||||
with pytest.raises(openai.InternalServerError):
|
||||
await run_client_embeddings(
|
||||
client,
|
||||
MODEL_NAME,
|
||||
sentences_1,
|
||||
)
|
||||
return
|
||||
|
||||
embeddings_list_1 = await run_client_embeddings(
|
||||
client,
|
||||
MODEL_NAME,
|
||||
@@ -117,7 +129,14 @@ def compute_lexical_matching_score(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
|
||||
async def test_bge_m3_api_server_sparse_embedding(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "token_classify":
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
await sparse_embeddings(client, sentences_1)
|
||||
return
|
||||
|
||||
embeddings_1 = await sparse_embeddings(client, sentences_1)
|
||||
embeddings_2 = await sparse_embeddings(client, sentences_2)
|
||||
|
||||
@@ -137,9 +156,11 @@ async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_sparse_embedding_corner_case(
|
||||
client: openai.AsyncOpenAI,
|
||||
):
|
||||
async def test_bge_m3_api_server_sparse_embedding_corner_case(server, pooling_task):
|
||||
if pooling_task != "token_classify":
|
||||
return
|
||||
|
||||
client = server.get_async_client()
|
||||
embeddings = await sparse_embeddings(client, ["Hi"])
|
||||
assert len(embeddings) == 1
|
||||
assert 2673 in embeddings[0]
|
||||
@@ -155,7 +176,18 @@ def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bge_m3_api_server_multi_vector(client: openai.AsyncOpenAI):
|
||||
async def test_bge_m3_api_server_multi_vector(server, pooling_task):
|
||||
client = server.get_async_client()
|
||||
|
||||
if pooling_task != "token_embed":
|
||||
with pytest.raises(openai.BadRequestError):
|
||||
await client.post(
|
||||
"../pooling",
|
||||
body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
|
||||
cast_to=httpx.Response,
|
||||
)
|
||||
return
|
||||
|
||||
result_1 = await client.post(
|
||||
"../pooling",
|
||||
body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
|
||||
|
||||
@@ -4,6 +4,7 @@ import pytest
|
||||
import torch
|
||||
|
||||
from vllm import TokensPrompt
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -20,6 +21,7 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
|
||||
max_model_len=128,
|
||||
enforce_eager=True,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
enable_prefix_caching=True,
|
||||
) as vllm_model:
|
||||
pooling_outputs = vllm_model.llm.encode(
|
||||
@@ -44,14 +46,3 @@ def test_extract_hidden_states(hf_runner, vllm_runner, model: str):
|
||||
assert len(output.prompt_token_ids) == n
|
||||
assert len(output.outputs.data) == n
|
||||
assert output.num_cached_tokens == 0
|
||||
|
||||
# skip_reading_prefix_cache can still write to cache
|
||||
# to accelerate following requests
|
||||
pooling_outputs = vllm_model.llm.encode(
|
||||
[TokensPrompt(prompt_token_ids=t) for t in token_prompts],
|
||||
pooling_task="embed",
|
||||
)
|
||||
|
||||
for n, output in zip(n_prompt_tokens, pooling_outputs):
|
||||
assert len(output.prompt_token_ids) == n
|
||||
assert output.num_cached_tokens > 0
|
||||
|
||||
@@ -5,6 +5,7 @@ import torch
|
||||
from transformers import AutoModel
|
||||
|
||||
from tests.models.utils import check_embeddings_close
|
||||
from vllm.config import PoolerConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -17,6 +18,7 @@ def test_embed_models(hf_runner, vllm_runner, example_prompts, model: str, dtype
|
||||
with vllm_runner(
|
||||
model,
|
||||
runner="pooling",
|
||||
pooler_config=PoolerConfig(task="token_embed"),
|
||||
max_model_len=None,
|
||||
) as vllm_model:
|
||||
vllm_outputs = vllm_model.token_embed(example_prompts)
|
||||
|
||||
@@ -146,7 +146,7 @@ def test_multi_vector_retrieval_models_using_normalize(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=False),
|
||||
pooler_config=PoolerConfig(use_activation=False, task="token_embed"),
|
||||
) as vllm_model:
|
||||
wo_normalize = vllm_model.token_embed(example_prompts)
|
||||
|
||||
@@ -154,7 +154,7 @@ def test_multi_vector_retrieval_models_using_normalize(
|
||||
model,
|
||||
max_model_len=512,
|
||||
dtype=dtype,
|
||||
pooler_config=PoolerConfig(use_activation=True),
|
||||
pooler_config=PoolerConfig(use_activation=True, task="token_embed"),
|
||||
) as vllm_model:
|
||||
w_normalize = vllm_model.token_embed(example_prompts)
|
||||
|
||||
|
||||
@@ -594,6 +594,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
|
||||
"MiMoV2FlashForCausalLM": _HfExamplesInfo(
|
||||
"XiaomiMiMo/MiMo-V2-Flash", trust_remote_code=True
|
||||
),
|
||||
"MiMoV2ProForCausalLM": _HfExamplesInfo(
|
||||
"XiaomiMiMo/MiMo-V2.5-Pro", trust_remote_code=True, is_available_online=False
|
||||
),
|
||||
"Dots1ForCausalLM": _HfExamplesInfo("rednote-hilab/dots.llm1.inst"),
|
||||
}
|
||||
|
||||
@@ -961,6 +964,18 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
"PerceptronAI/Isaac-0.1",
|
||||
trust_remote_code=True,
|
||||
extras={"0.2-2B-Preview": "PerceptronAI/Isaac-0.2-2B-Preview"},
|
||||
max_transformers_version="4.57",
|
||||
transformers_version_reason={
|
||||
"vllm": (
|
||||
"Custom Isaac code is not compatible with Transformers v5. "
|
||||
"The model should be upstreamed to Transformers for "
|
||||
"long-term support."
|
||||
),
|
||||
"hf": (
|
||||
"Isaac's remote model and processor code import or configure "
|
||||
"APIs that changed in Transformers v5."
|
||||
),
|
||||
},
|
||||
),
|
||||
"InternS1ForConditionalGeneration": _HfExamplesInfo(
|
||||
"internlm/Intern-S1",
|
||||
@@ -1057,6 +1072,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
"MiDashengLMModel": _HfExamplesInfo(
|
||||
"mispeech/midashenglm-7b", trust_remote_code=True
|
||||
),
|
||||
"MiMoV2OmniForCausalLM": _HfExamplesInfo(
|
||||
"XiaomiMiMo/MiMo-V2.5-Omni", trust_remote_code=True, is_available_online=False
|
||||
),
|
||||
"MiniCPMO": _HfExamplesInfo(
|
||||
"openbmb/MiniCPM-o-2_6",
|
||||
trust_remote_code=True,
|
||||
@@ -1540,6 +1558,18 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
|
||||
trust_remote_code=True,
|
||||
speculative_model="XiaomiMiMo/MiMo-7B-RL",
|
||||
),
|
||||
"MiMoV2MTPModel": _HfExamplesInfo(
|
||||
"XiaomiMiMo/MiMo-V2.5-Pro",
|
||||
trust_remote_code=True,
|
||||
speculative_model="XiaomiMiMo/MiMo-V2.5-Pro",
|
||||
is_available_online=False,
|
||||
),
|
||||
"MiMoV2OmniMTPModel": _HfExamplesInfo(
|
||||
"XiaomiMiMo/MiMo-V2.5-Omni",
|
||||
trust_remote_code=True,
|
||||
speculative_model="XiaomiMiMo/MiMo-V2.5-Omni",
|
||||
is_available_online=False,
|
||||
),
|
||||
"NemotronHMTPModel": _HfExamplesInfo(
|
||||
"nvidia/Nemotron-Super-Placeholder",
|
||||
speculative_model="nvidia/Nemotron-Super-Placeholder",
|
||||
|
||||
@@ -5,6 +5,8 @@ Unit tests for MultiModalRegistry.supports_multimodal_inputs and
|
||||
Qwen2.5-VL visual component loading behavior.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
@@ -32,3 +34,17 @@ def test_supports_multimodal_inputs(model_id, limit_mm_per_prompt, expected):
|
||||
limit_mm_per_prompt=limit_mm_per_prompt,
|
||||
)
|
||||
assert MULTIMODAL_REGISTRY.supports_multimodal_inputs(ctx.model_config) is expected
|
||||
|
||||
|
||||
def test_create_processor_error_uses_served_model_name():
|
||||
model_config = SimpleNamespace(
|
||||
is_multimodal_model=False,
|
||||
model="/path/to/model/weights",
|
||||
served_model_name="friendly-model-name",
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="friendly-model-name is not a multimodal model",
|
||||
):
|
||||
MULTIMODAL_REGISTRY.create_processor(model_config)
|
||||
|
||||
@@ -0,0 +1,185 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for Cutlass W4A16 (Machete) kernel on Hopper.
|
||||
|
||||
Verifies that W4A16 quantized models loaded through vllm select the
|
||||
MacheteLinearKernel on sm_90 GPUs, that weights are correctly repacked,
|
||||
and that inference produces valid output.
|
||||
|
||||
Run `pytest tests/quantization/test_cutlass_w4a16.py`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.has_device_capability(90):
|
||||
pytest.skip(
|
||||
"Machete W4A16 requires Hopper (sm_90).",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
MPLinearLayerConfig,
|
||||
choose_mp_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.mixed_precision import (
|
||||
MacheteLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
|
||||
CompressedTensorsLinearMethod,
|
||||
CompressedTensorsWNA16,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
|
||||
@pytest.fixture(scope="function", autouse=True)
|
||||
def enable_pickle(monkeypatch):
|
||||
"""`LLM.apply_model` requires pickling a function."""
|
||||
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"act_type,weight_type,group_size,zero_points",
|
||||
[
|
||||
(torch.float16, scalar_types.uint4b8, 128, False),
|
||||
(torch.bfloat16, scalar_types.uint4b8, 128, False),
|
||||
(torch.float16, scalar_types.uint4, 128, True),
|
||||
(torch.float16, scalar_types.uint4b8, -1, False),
|
||||
],
|
||||
ids=[
|
||||
"fp16-gptq-g128",
|
||||
"bf16-gptq-g128",
|
||||
"fp16-awq-g128",
|
||||
"fp16-channelwise",
|
||||
],
|
||||
)
|
||||
def test_machete_kernel_selected(act_type, weight_type, group_size, zero_points):
|
||||
"""Verify choose_mp_linear_kernel picks MacheteLinearKernel."""
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(4096, 4096),
|
||||
partition_weight_shape=(4096, 4096),
|
||||
act_type=act_type,
|
||||
weight_type=weight_type,
|
||||
group_size=group_size,
|
||||
zero_points=zero_points,
|
||||
has_g_idx=False,
|
||||
)
|
||||
kernel = choose_mp_linear_kernel(config)
|
||||
assert kernel is MacheteLinearKernel, (
|
||||
f"Expected MacheteLinearKernel, got {kernel.__name__}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"full_shape,part_shape,weight_type,group_size,has_g_idx,expected_reason",
|
||||
[
|
||||
((4096, 4096), (2048, 4096), scalar_types.uint4b8, 128, True, "Act reordering"),
|
||||
(
|
||||
(4096, 4096),
|
||||
(4096, 4096),
|
||||
scalar_types.float6_e3m2f,
|
||||
128,
|
||||
False,
|
||||
"Quant type",
|
||||
),
|
||||
((4096, 4096), (4096, 4096), scalar_types.uint4b8, 32, False, "Group size"),
|
||||
],
|
||||
ids=["partitioned-g_idx", "unsupported-quant-type", "unsupported-group-size"],
|
||||
)
|
||||
def test_machete_rejects_invalid_config(
|
||||
full_shape, part_shape, weight_type, group_size, has_g_idx, expected_reason
|
||||
):
|
||||
"""Verify Machete rejects unsupported configurations."""
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=full_shape,
|
||||
partition_weight_shape=part_shape,
|
||||
act_type=torch.float16,
|
||||
weight_type=weight_type,
|
||||
group_size=group_size,
|
||||
zero_points=False,
|
||||
has_g_idx=has_g_idx,
|
||||
)
|
||||
can_impl, reason = MacheteLinearKernel.can_implement(config)
|
||||
assert not can_impl
|
||||
assert expected_reason in reason
|
||||
|
||||
|
||||
def test_kernel_selection_with_disabled_machete(monkeypatch):
|
||||
"""Verify kernel selection falls back when Machete is disabled."""
|
||||
monkeypatch.setattr("vllm.envs.VLLM_DISABLED_KERNELS", ["MacheteLinearKernel"])
|
||||
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(4096, 4096),
|
||||
partition_weight_shape=(4096, 4096),
|
||||
act_type=torch.float16,
|
||||
weight_type=scalar_types.uint4b8,
|
||||
group_size=128,
|
||||
zero_points=False,
|
||||
has_g_idx=False,
|
||||
)
|
||||
kernel = choose_mp_linear_kernel(config)
|
||||
assert kernel is not MacheteLinearKernel, "MacheteLinearKernel should be disabled"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[
|
||||
"nm-testing/tinyllama-oneshot-w4a16-channel-v2",
|
||||
"nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder",
|
||||
],
|
||||
)
|
||||
def test_w4a16_machete_e2e(vllm_runner, model_name):
|
||||
"""Load a W4A16 model, verify Machete kernel is used, and generate."""
|
||||
with vllm_runner(model_name, enforce_eager=True, gpu_memory_utilization=0.5) as llm:
|
||||
|
||||
def check_model(model):
|
||||
layer = model.model.layers[0]
|
||||
qkv_proj = layer.self_attn.qkv_proj
|
||||
|
||||
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
|
||||
assert isinstance(qkv_proj.scheme, CompressedTensorsWNA16)
|
||||
assert isinstance(qkv_proj.scheme.kernel, MacheteLinearKernel), (
|
||||
f"Expected MacheteLinearKernel on Hopper, "
|
||||
f"got {type(qkv_proj.scheme.kernel).__name__}"
|
||||
)
|
||||
|
||||
assert hasattr(qkv_proj, "weight_packed")
|
||||
assert hasattr(qkv_proj, "weight_scale")
|
||||
assert qkv_proj.weight_packed.dtype == torch.int32
|
||||
|
||||
llm.apply_model(check_model)
|
||||
|
||||
output = llm.generate_greedy("Hello my name is", max_tokens=10)
|
||||
assert output
|
||||
assert len(output[0][1]) > 0
|
||||
|
||||
|
||||
def test_w4a16_machete_bfloat16_deterministic(vllm_runner):
|
||||
"""Verify Machete works with bf16 activations and is deterministic."""
|
||||
model_name = "nm-testing/tinyllama-oneshot-w4a16-channel-v2"
|
||||
prompt = "The capital of France is"
|
||||
|
||||
with vllm_runner(
|
||||
model_name,
|
||||
enforce_eager=True,
|
||||
dtype="bfloat16",
|
||||
gpu_memory_utilization=0.5,
|
||||
) as llm:
|
||||
|
||||
def check_kernel_type(model):
|
||||
layer = model.model.layers[0]
|
||||
scheme = layer.self_attn.qkv_proj.scheme
|
||||
assert isinstance(scheme.kernel, MacheteLinearKernel), (
|
||||
f"Expected MacheteLinearKernel with bf16, "
|
||||
f"got {type(scheme.kernel).__name__}"
|
||||
)
|
||||
|
||||
llm.apply_model(check_kernel_type)
|
||||
|
||||
out1 = llm.generate_greedy(prompt, max_tokens=10)
|
||||
out2 = llm.generate_greedy(prompt, max_tokens=10)
|
||||
assert out1[0][1] == out2[0][1], (
|
||||
f"Non-deterministic: '{out1[0][1]}' vs '{out2[0][1]}'"
|
||||
)
|
||||
@@ -188,30 +188,6 @@ class TestExtractToolCalls:
|
||||
"location": "NYC"
|
||||
}
|
||||
|
||||
def test_type_conversion_in_non_streaming(self):
|
||||
"""Non-streaming extraction must convert params using the tool schema."""
|
||||
tool = ChatCompletionToolsParam(
|
||||
function=FunctionDefinition(
|
||||
name="toggle",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"enabled": {"type": "boolean"},
|
||||
"count": {"type": "integer"},
|
||||
},
|
||||
},
|
||||
),
|
||||
)
|
||||
parser = make_parser(tools=[tool])
|
||||
model_output = build_tool_call("toggle", {"enabled": "true", "count": "42"})
|
||||
result = parser.extract_tool_calls(model_output, None)
|
||||
assert result.tools_called
|
||||
assert len(result.tool_calls) == 1
|
||||
args = json.loads(result.tool_calls[0].function.arguments)
|
||||
assert args == {"enabled": True, "count": 42}
|
||||
assert isinstance(args["enabled"], bool)
|
||||
assert isinstance(args["count"], int)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: extract_tool_calls_streaming
|
||||
|
||||
@@ -2074,54 +2074,6 @@ def test_auto_fit_max_model_len_not_triggered():
|
||||
assert vllm_config.model_config.max_model_len == 16
|
||||
|
||||
|
||||
def test_auto_fit_max_model_len_respects_num_gpu_blocks_override():
|
||||
"""Auto-fit must size max_model_len against the override-clamped pool, not
|
||||
the raw `available_memory`. Without this, auto-fit could pick a
|
||||
max_model_len that no longer fits once `num_gpu_blocks_override` is applied.
|
||||
"""
|
||||
model_config = ModelConfig(max_model_len=16384)
|
||||
model_config.original_max_model_len = -1 # request auto-fit
|
||||
vllm_config = VllmConfig(model_config=model_config)
|
||||
# Cap the cache to 32 blocks regardless of available memory.
|
||||
vllm_config.cache_config.num_gpu_blocks_override = 32
|
||||
|
||||
mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
|
||||
kv_cache_specs = {
|
||||
"layer_1": new_kv_cache_spec(), # block_size=16
|
||||
"layer_2": new_kv_cache_spec(),
|
||||
}
|
||||
# Plenty of raw memory (1024 blocks per layer would fit max_model_len=16384).
|
||||
large_available_memory = mem_per_block_per_layer * 2 * 1024
|
||||
|
||||
get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
|
||||
|
||||
# 32 blocks * block_size 16 = 512 token slots, so max_model_len must
|
||||
# auto-fit at or below that.
|
||||
assert 0 < vllm_config.model_config.max_model_len <= 32 * 16
|
||||
|
||||
|
||||
def test_check_enough_kv_cache_memory_respects_num_gpu_blocks_override():
|
||||
"""Admission check must use the override-clamped pool size, not raw
|
||||
`available_memory`. Without this, startup could accept a max_model_len
|
||||
that does not actually fit in `num_gpu_blocks_override` blocks.
|
||||
"""
|
||||
model_config = ModelConfig(max_model_len=16384)
|
||||
vllm_config = VllmConfig(model_config=model_config)
|
||||
# 32 blocks is far too small for max_model_len=16384 (would need 1024).
|
||||
vllm_config.cache_config.num_gpu_blocks_override = 32
|
||||
|
||||
mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
|
||||
kv_cache_specs = {
|
||||
"layer_1": new_kv_cache_spec(),
|
||||
"layer_2": new_kv_cache_spec(),
|
||||
}
|
||||
# Plenty of raw memory: a bytes-only check against this would pass.
|
||||
large_available_memory = mem_per_block_per_layer * 2 * 1024
|
||||
|
||||
with pytest.raises(ValueError, match="max seq len"):
|
||||
get_kv_cache_configs(vllm_config, [kv_cache_specs], [large_available_memory])
|
||||
|
||||
|
||||
def test_unify_hybrid_kv_cache_specs():
|
||||
# 1. has_full_attention and has_sliding_window
|
||||
before_spec_1 = new_kv_cache_spec()
|
||||
|
||||
@@ -324,13 +324,10 @@ def run_test(
|
||||
):
|
||||
spec_decoding = spec_config is not None
|
||||
cache_arg: dict[str, Any] = (
|
||||
# Force preemptions: with 32 blocks the cache holds at most a single
|
||||
# max-length request, so the ~34 concurrent prompts contend and trigger
|
||||
# preemption. (Prompts here are << max_model_len, so dropping
|
||||
# max_model_len from 4096 to 512 doesn't change generation behavior.)
|
||||
dict(num_gpu_blocks_override=32, max_model_len=512)
|
||||
# Force preemptions
|
||||
dict(num_gpu_blocks_override=32)
|
||||
if test_preemption
|
||||
else dict(gpu_memory_utilization=0.9, max_model_len=4096)
|
||||
else dict(gpu_memory_utilization=0.9)
|
||||
)
|
||||
spec_mml = (spec_config or {}).get("max_model_len")
|
||||
spec_method = (spec_config or {}).get("method", "none")
|
||||
@@ -346,6 +343,7 @@ def run_test(
|
||||
|
||||
with VllmRunner(
|
||||
model,
|
||||
max_model_len=4096,
|
||||
enable_chunked_prefill=test_prefill_chunking,
|
||||
# Force prefill chunking
|
||||
max_num_batched_tokens=48 if test_prefill_chunking else None,
|
||||
|
||||
@@ -8,16 +8,11 @@ from unittest.mock import Mock
|
||||
import pytest
|
||||
|
||||
from vllm.config import ModelConfig, SchedulerConfig, VllmConfig
|
||||
from vllm.reasoning import ReasoningParser
|
||||
from vllm.v1.request import Request
|
||||
from vllm.v1.structured_output import StructuredOutputManager
|
||||
|
||||
|
||||
class MockReasoner:
|
||||
def __init__(self, tokenizer):
|
||||
self.is_reasoning_end = Mock(return_value=False)
|
||||
self.is_reasoning_end_streaming = Mock(return_value=False)
|
||||
|
||||
|
||||
class TestReasoningStructuredOutput:
|
||||
"""Test reasoning-aware structured output functionality."""
|
||||
|
||||
@@ -55,6 +50,13 @@ class TestReasoningStructuredOutput:
|
||||
config.speculative_config = None
|
||||
return config
|
||||
|
||||
@pytest.fixture
|
||||
def mock_reasoning_parser(self):
|
||||
"""Create a mock ReasoningParser."""
|
||||
parser = Mock(spec=ReasoningParser)
|
||||
parser.is_reasoning_end = Mock(return_value=False)
|
||||
return parser
|
||||
|
||||
@pytest.fixture
|
||||
def mock_request_with_structured_output(self):
|
||||
"""Create a mock request with structured output."""
|
||||
@@ -62,8 +64,6 @@ class TestReasoningStructuredOutput:
|
||||
request.structured_output_request = Mock()
|
||||
request.structured_output_request.reasoning_ended = None
|
||||
request.structured_output_request.grammar = Mock()
|
||||
request.structured_output_request.reasoning_parser_kwargs = None
|
||||
request.structured_output_request.reasoner = None
|
||||
request.structured_output_request.grammar.is_terminated = Mock(
|
||||
return_value=False
|
||||
)
|
||||
@@ -74,13 +74,6 @@ class TestReasoningStructuredOutput:
|
||||
request.num_output_placeholders = 0
|
||||
return request
|
||||
|
||||
@pytest.fixture
|
||||
def manager_with_reasoner(self, mock_vllm_config):
|
||||
manager = StructuredOutputManager(mock_vllm_config)
|
||||
manager.reasoner_cls = MockReasoner
|
||||
manager.tokenizer = Mock()
|
||||
return manager
|
||||
|
||||
def test_should_fill_bitmask_with_enable_in_reasoning(
|
||||
self, mock_vllm_config, mock_request_with_structured_output
|
||||
):
|
||||
@@ -96,17 +89,22 @@ class TestReasoningStructuredOutput:
|
||||
|
||||
def test_should_fill_bitmask_without_enable_in_reasoning(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_vllm_config,
|
||||
mock_request_with_structured_output,
|
||||
mock_reasoning_parser,
|
||||
):
|
||||
"""Test should_fill_bitmask when enable_in_reasoning is False."""
|
||||
# Keep enable_in_reasoning as False (default)
|
||||
config = manager_with_reasoner.vllm_config.structured_outputs_config
|
||||
config = mock_vllm_config.structured_outputs_config
|
||||
assert config.enable_in_reasoning is False
|
||||
|
||||
result = manager_with_reasoner.should_fill_bitmask(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
manager = StructuredOutputManager(mock_vllm_config)
|
||||
manager.reasoner = mock_reasoning_parser
|
||||
|
||||
# Mock reasoning not ended
|
||||
mock_reasoning_parser.is_reasoning_end.return_value = False
|
||||
|
||||
result = manager.should_fill_bitmask(mock_request_with_structured_output)
|
||||
|
||||
# Should set reasoning_ended and return its value
|
||||
assert (
|
||||
@@ -120,92 +118,68 @@ class TestReasoningStructuredOutput:
|
||||
):
|
||||
"""Test should_fill_bitmask when no reasoner is configured."""
|
||||
manager = StructuredOutputManager(mock_vllm_config)
|
||||
manager.reasoner = None
|
||||
|
||||
result = manager.should_fill_bitmask(mock_request_with_structured_output)
|
||||
|
||||
# Should default to True when no reasoner
|
||||
assert result is True
|
||||
|
||||
def test_should_fill_bitmask_uses_request_reasoning_parser_kwargs(
|
||||
self, mock_vllm_config, mock_request_with_structured_output
|
||||
):
|
||||
"""Test request-level parser kwargs override the default reasoner."""
|
||||
|
||||
class KwargReasoner:
|
||||
def __init__(self, tokenizer, chat_template_kwargs=None):
|
||||
self.chat_template_kwargs = chat_template_kwargs or {}
|
||||
|
||||
def is_reasoning_end(self, input_ids):
|
||||
return not self.chat_template_kwargs.get("enable_thinking", False)
|
||||
|
||||
manager = StructuredOutputManager(mock_vllm_config)
|
||||
manager.reasoner_cls = KwargReasoner
|
||||
manager.tokenizer = Mock()
|
||||
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoning_parser_kwargs = {
|
||||
"chat_template_kwargs": {"enable_thinking": True}
|
||||
}
|
||||
|
||||
result = manager.should_fill_bitmask(mock_request_with_structured_output)
|
||||
|
||||
assert result is False
|
||||
assert (
|
||||
mock_request_with_structured_output.structured_output_request.reasoner
|
||||
is not None
|
||||
)
|
||||
|
||||
def test_should_advance_with_enable_in_reasoning(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_vllm_config,
|
||||
mock_request_with_structured_output,
|
||||
mock_reasoning_parser,
|
||||
):
|
||||
"""Test should_advance when enable_in_reasoning is True."""
|
||||
# Enable enable_in_reasoning
|
||||
manager_with_reasoner.enable_in_reasoning = True
|
||||
mock_vllm_config.structured_outputs_config.enable_in_reasoning = True
|
||||
|
||||
manager = StructuredOutputManager(mock_vllm_config)
|
||||
manager.reasoner = mock_reasoning_parser
|
||||
|
||||
# Should always return True when enable_in_reasoning is enabled
|
||||
result = manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
result = manager.should_advance(mock_request_with_structured_output)
|
||||
assert result is True
|
||||
|
||||
def test_should_advance_reasoning_not_ended(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_vllm_config,
|
||||
mock_request_with_structured_output,
|
||||
mock_reasoning_parser,
|
||||
):
|
||||
"""Test should_advance when reasoning has not ended."""
|
||||
manager = StructuredOutputManager(mock_vllm_config)
|
||||
manager.reasoner = mock_reasoning_parser
|
||||
|
||||
# Set reasoning as not ended
|
||||
(
|
||||
mock_request_with_structured_output.structured_output_request
|
||||
).reasoning_ended = False
|
||||
mock_reasoning_parser.is_reasoning_end.return_value = False
|
||||
|
||||
result = manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
result = manager.should_advance(mock_request_with_structured_output)
|
||||
|
||||
# Should return False since reasoning hasn't ended
|
||||
assert result is False
|
||||
|
||||
def test_should_advance_reasoning_just_ended(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_vllm_config,
|
||||
mock_request_with_structured_output,
|
||||
mock_reasoning_parser,
|
||||
):
|
||||
"""Test should_advance when reasoning ends in current step."""
|
||||
manager = StructuredOutputManager(mock_vllm_config)
|
||||
manager.reasoner = mock_reasoning_parser
|
||||
|
||||
# Set reasoning as not ended initially, but ends in this step
|
||||
(
|
||||
mock_request_with_structured_output.structured_output_request
|
||||
).reasoning_ended = False
|
||||
reasoner = MockReasoner(tokenizer=Mock())
|
||||
reasoner.is_reasoning_end_streaming.return_value = True
|
||||
structured_req = mock_request_with_structured_output.structured_output_request
|
||||
structured_req.reasoner = reasoner
|
||||
mock_reasoning_parser.is_reasoning_end.return_value = True
|
||||
|
||||
result = manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
result = manager.should_advance(mock_request_with_structured_output)
|
||||
|
||||
# Should set reasoning_ended to True but return False for this step
|
||||
assert (
|
||||
@@ -216,18 +190,20 @@ class TestReasoningStructuredOutput:
|
||||
|
||||
def test_should_advance_reasoning_already_ended(
|
||||
self,
|
||||
manager_with_reasoner,
|
||||
mock_vllm_config,
|
||||
mock_request_with_structured_output,
|
||||
mock_reasoning_parser,
|
||||
):
|
||||
"""Test should_advance when reasoning has already ended."""
|
||||
manager = StructuredOutputManager(mock_vllm_config)
|
||||
manager.reasoner = mock_reasoning_parser
|
||||
|
||||
# Set reasoning as already ended
|
||||
(
|
||||
mock_request_with_structured_output.structured_output_request
|
||||
).reasoning_ended = True
|
||||
|
||||
result = manager_with_reasoner.should_advance(
|
||||
mock_request_with_structured_output
|
||||
)
|
||||
result = manager.should_advance(mock_request_with_structured_output)
|
||||
|
||||
# Should return True since reasoning has ended
|
||||
assert result is True
|
||||
|
||||
+13
-1
@@ -326,6 +326,10 @@ class ModelConfig:
|
||||
mm_encoder_only: InitVar[bool | None] = None
|
||||
mm_encoder_tp_mode: InitVar[MMEncoderTPMode | None] = None
|
||||
mm_encoder_attn_backend: InitVar[AttentionBackendEnum | str | None] = None
|
||||
mm_encoder_attn_dtype: InitVar[str | None] = None
|
||||
mm_encoder_fp8_scale_path: InitVar[str | None] = None
|
||||
mm_encoder_fp8_scale_save_path: InitVar[str | None] = None
|
||||
mm_encoder_fp8_scale_save_margin: InitVar[float | None] = None
|
||||
interleave_mm_strings: InitVar[bool | None] = None
|
||||
skip_mm_profiling: InitVar[bool | None] = None
|
||||
video_pruning_rate: InitVar[float | None] = None
|
||||
@@ -447,6 +451,10 @@ class ModelConfig:
|
||||
mm_encoder_only: bool | None,
|
||||
mm_encoder_tp_mode: MMEncoderTPMode | None,
|
||||
mm_encoder_attn_backend: AttentionBackendEnum | str | None,
|
||||
mm_encoder_attn_dtype: str | None,
|
||||
mm_encoder_fp8_scale_path: str | None,
|
||||
mm_encoder_fp8_scale_save_path: str | None,
|
||||
mm_encoder_fp8_scale_save_margin: float | None,
|
||||
interleave_mm_strings: bool | None,
|
||||
skip_mm_profiling: bool | None,
|
||||
video_pruning_rate: float | None,
|
||||
@@ -513,6 +521,7 @@ class ModelConfig:
|
||||
if dict_overrides:
|
||||
self._apply_dict_overrides(hf_config, dict_overrides)
|
||||
self.hf_text_config = get_hf_text_config(self.hf_config)
|
||||
self.model_arch_config = self.get_model_arch_config()
|
||||
self.attention_chunk_size = getattr(
|
||||
self.hf_text_config, "attention_chunk_size", None
|
||||
)
|
||||
@@ -520,7 +529,6 @@ class ModelConfig:
|
||||
self.hf_image_processor_config = get_hf_image_processor_config(
|
||||
self.model, hf_token=self.hf_token, revision=self.revision
|
||||
)
|
||||
self.model_arch_config = self.get_model_arch_config()
|
||||
|
||||
architectures = self.architectures
|
||||
registry = self.registry
|
||||
@@ -643,6 +651,10 @@ class ModelConfig:
|
||||
mm_encoder_only=mm_encoder_only,
|
||||
mm_encoder_tp_mode=mm_encoder_tp_mode,
|
||||
mm_encoder_attn_backend=mm_encoder_attn_backend,
|
||||
mm_encoder_attn_dtype=mm_encoder_attn_dtype,
|
||||
mm_encoder_fp8_scale_path=mm_encoder_fp8_scale_path,
|
||||
mm_encoder_fp8_scale_save_path=mm_encoder_fp8_scale_save_path,
|
||||
mm_encoder_fp8_scale_save_margin=mm_encoder_fp8_scale_save_margin,
|
||||
interleave_mm_strings=interleave_mm_strings,
|
||||
skip_mm_profiling=skip_mm_profiling,
|
||||
video_pruning_rate=video_pruning_rate,
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, TypeAlias, TypedDict, final
|
||||
|
||||
from pydantic import ConfigDict, Field, field_validator, model_validator
|
||||
@@ -158,6 +159,24 @@ class MultiModalConfig:
|
||||
"""Optional override for the multi-modal encoder attention backend when
|
||||
using vision transformers. Accepts any value from
|
||||
`vllm.v1.attention.backends.registry.AttentionBackendEnum` (e.g. `FLASH_ATTN`)."""
|
||||
mm_encoder_attn_dtype: Literal["fp8"] | None = None
|
||||
"""Optional dtype override for ViT encoder attention. Set to `"fp8"` to
|
||||
enable FP8 quantization via the FlashInfer cuDNN backend. When set to
|
||||
`"fp8"` without a scale file, dynamic scaling is used automatically.
|
||||
See docs/features/quantization/fp8_vit_attn.md for details."""
|
||||
mm_encoder_fp8_scale_path: str | None = None
|
||||
"""Path to a JSON file containing per-layer FP8 Q/K/V scales for ViT
|
||||
encoder attention. When provided (with `mm_encoder_attn_dtype="fp8"`),
|
||||
static scaling is used. When omitted, dynamic scaling is used."""
|
||||
mm_encoder_fp8_scale_save_path: str | None = None
|
||||
"""When set with dynamic FP8 scaling (`mm_encoder_attn_dtype="fp8"`
|
||||
and no `mm_encoder_fp8_scale_path`), saves the calibrated scales to
|
||||
this file after the amax history buffer is full. The saved file can
|
||||
then be used as `mm_encoder_fp8_scale_path` in subsequent runs."""
|
||||
mm_encoder_fp8_scale_save_margin: float = Field(default=1.5, gt=0.0)
|
||||
"""Safety margin multiplied onto scales when auto-saving. A value > 1
|
||||
leaves headroom so that inputs with larger activations than the
|
||||
calibration set do not overflow FP8 range. Default 1.5."""
|
||||
interleave_mm_strings: bool = False
|
||||
"""Enable fully interleaved support for multimodal prompts, while using
|
||||
--chat-template-content-format=string."""
|
||||
@@ -233,6 +252,36 @@ class MultiModalConfig:
|
||||
"'mm_shm_cache_max_object_size_mb' should only be set when "
|
||||
"'mm_processor_cache_type' is 'shm'."
|
||||
)
|
||||
# Validate FP8 scale path combinations.
|
||||
if self.mm_encoder_attn_dtype != "fp8" and (
|
||||
self.mm_encoder_fp8_scale_path is not None
|
||||
or self.mm_encoder_fp8_scale_save_path is not None
|
||||
):
|
||||
raise ValueError(
|
||||
"'mm_encoder_fp8_scale_path' and "
|
||||
"'mm_encoder_fp8_scale_save_path' require "
|
||||
"'mm_encoder_attn_dtype' to be 'fp8'."
|
||||
)
|
||||
if (
|
||||
self.mm_encoder_fp8_scale_path is not None
|
||||
and self.mm_encoder_fp8_scale_save_path is not None
|
||||
):
|
||||
raise ValueError(
|
||||
"'mm_encoder_fp8_scale_save_path' cannot be used with "
|
||||
"'mm_encoder_fp8_scale_path' (saving requires dynamic scaling)."
|
||||
)
|
||||
|
||||
# Validate file paths exist.
|
||||
if self.mm_encoder_fp8_scale_path is not None:
|
||||
scale_path = Path(self.mm_encoder_fp8_scale_path)
|
||||
if not scale_path.is_file():
|
||||
raise FileNotFoundError(f"FP8 scale file not found: {scale_path}")
|
||||
if self.mm_encoder_fp8_scale_save_path is not None:
|
||||
save_parent = Path(self.mm_encoder_fp8_scale_save_path).parent
|
||||
if not save_parent.is_dir():
|
||||
raise FileNotFoundError(
|
||||
f"Parent directory for FP8 scale save path not found: {save_parent}"
|
||||
)
|
||||
return self
|
||||
|
||||
def compute_hash(self) -> str:
|
||||
@@ -252,6 +301,8 @@ class MultiModalConfig:
|
||||
if self.mm_encoder_attn_backend is not None
|
||||
else None,
|
||||
self.mm_encoder_tp_mode,
|
||||
self.mm_encoder_attn_dtype,
|
||||
self.mm_encoder_fp8_scale_path,
|
||||
]
|
||||
hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
|
||||
return hash_str
|
||||
|
||||
@@ -34,6 +34,7 @@ logger = init_logger(__name__)
|
||||
MTPModelTypes = Literal[
|
||||
"deepseek_mtp",
|
||||
"mimo_mtp",
|
||||
"mimo_v2_mtp",
|
||||
"glm4_moe_mtp",
|
||||
"glm4_moe_lite_mtp",
|
||||
"glm_ocr_mtp",
|
||||
@@ -63,7 +64,8 @@ SpeculativeMethod = Literal[
|
||||
EagleModelTypes,
|
||||
NgramGPUTypes,
|
||||
]
|
||||
RejectionSampleMethod = Literal["strict", "probabilistic", "synthetic"]
|
||||
RejectionSampleMethod = Literal["standard", "synthetic"]
|
||||
DraftSampleMethod = Literal["greedy", "gumbel"]
|
||||
|
||||
|
||||
@config
|
||||
@@ -183,11 +185,11 @@ class SpeculativeConfig:
|
||||
"""Load config for the draft model. If not specified, will use the load
|
||||
config from the target model."""
|
||||
|
||||
rejection_sample_method: RejectionSampleMethod = "strict"
|
||||
"""Whether to use strict (target and draft sampled tokens match exactly)
|
||||
or probabilistic rejection sampling. Both respect the target model
|
||||
distribution, but the latter yields a higher acceptance rate at the cost
|
||||
of more memory to cache draft logits."""
|
||||
rejection_sample_method: RejectionSampleMethod = "standard"
|
||||
"""The rejection sampling method to use. 'standard' uses probabilistic
|
||||
rejection sampling (with or without cached draft logits, controlled by
|
||||
draft_sample_method). 'synthetic' accepts draft tokens with a decaying
|
||||
probability calibrated to synthetic_acceptance_rate."""
|
||||
|
||||
synthetic_acceptance_rates: list[float] | None = None
|
||||
"""Per-position *unconditional* acceptance rates for synthetic rejection
|
||||
@@ -248,6 +250,14 @@ class SpeculativeConfig:
|
||||
)
|
||||
return SpeculativeConfig._acceptance_length_to_rates(length, n)
|
||||
|
||||
draft_sample_method: DraftSampleMethod = "greedy"
|
||||
"""How the draft model samples tokens. 'greedy' always picks the argmax
|
||||
token, and the draft probabilities are treated as one-hot during rejection
|
||||
sampling. 'gumbel' adds Gumbel noise for stochastic sampling, and the full
|
||||
draft logits are used for the probability ratio test during rejection
|
||||
sampling. This comes at the cost of additional GPU memory usage. This
|
||||
parameter currently only applies to Model Runner V2."""
|
||||
|
||||
def compute_hash(self) -> str:
|
||||
"""
|
||||
WARNING: Whenever a new field is added to this config,
|
||||
@@ -323,6 +333,48 @@ class SpeculativeConfig:
|
||||
}
|
||||
)
|
||||
|
||||
if (arch := hf_config.architectures[0]) in (
|
||||
"MiMoV2ProForCausalLM",
|
||||
"MiMoV2OmniForCausalLM",
|
||||
):
|
||||
from vllm.model_executor.models.mimo_v2_mtp import (
|
||||
_MIMO_V2_PRO_NUM_MTP_LAYERS,
|
||||
)
|
||||
|
||||
mtp_arch_maps = {
|
||||
"MiMoV2ProForCausalLM": "MiMoV2MTPModel",
|
||||
"MiMoV2OmniForCausalLM": "MiMoV2OmniMTPModel",
|
||||
}
|
||||
|
||||
hf_config.model_type = "mimo_v2_mtp"
|
||||
# vLLM currently supports only the first MiMo-V2 MTP layer.
|
||||
n_predict = _MIMO_V2_PRO_NUM_MTP_LAYERS
|
||||
hf_config.update(
|
||||
{
|
||||
"num_hidden_layers": 0,
|
||||
"n_predict": n_predict,
|
||||
"num_nextn_predict_layers": n_predict,
|
||||
"architectures": [mtp_arch_maps[arch]],
|
||||
}
|
||||
)
|
||||
|
||||
if hf_config.architectures[0] == "MiMoV2FlashForCausalLM":
|
||||
from vllm.model_executor.models.mimo_v2_mtp import (
|
||||
_MIMO_V2_FLASH_NUM_MTP_LAYERS,
|
||||
)
|
||||
|
||||
hf_config.model_type = "mimo_v2_mtp"
|
||||
# vLLM currently supports only the first MiMo-V2 MTP layer.
|
||||
n_predict = _MIMO_V2_FLASH_NUM_MTP_LAYERS
|
||||
hf_config.update(
|
||||
{
|
||||
"num_hidden_layers": 0,
|
||||
"n_predict": n_predict,
|
||||
"num_nextn_predict_layers": n_predict,
|
||||
"architectures": ["MiMoV2MTPModel"],
|
||||
}
|
||||
)
|
||||
|
||||
if hf_config.architectures[0] == "Glm4MoeForCausalLM":
|
||||
hf_config.model_type = "glm4_moe_mtp"
|
||||
n_predict = getattr(hf_config, "num_nextn_predict_layers", None)
|
||||
|
||||
@@ -1432,10 +1432,6 @@ class VllmConfig:
|
||||
cudagraph_capture_sizes = [1, 2, 4] + list(range(8, 256, 8)) + list(
|
||||
range(256, max_graph_size + 1, 16))
|
||||
|
||||
`max_num_batched_tokens` is also appended to the list if it fits
|
||||
within `max_cudagraph_capture_size`, so the max batch size is captured
|
||||
even when off-stride.
|
||||
|
||||
In the end, `vllm_config.compilation_config.cudagraph_capture_sizes`
|
||||
will be the final sizes to capture cudagraph (in ascending order).
|
||||
|
||||
@@ -1524,12 +1520,6 @@ class VllmConfig:
|
||||
cudagraph_capture_sizes += list(
|
||||
range(256, max_cudagraph_capture_size + 1, 16)
|
||||
)
|
||||
# ensure max_num_tokens is captured if within max capture size
|
||||
if (
|
||||
max_num_tokens <= max_cudagraph_capture_size
|
||||
and max_num_tokens not in cudagraph_capture_sizes
|
||||
):
|
||||
cudagraph_capture_sizes.append(max_num_tokens)
|
||||
# de-duplicate and sort the sizes
|
||||
cudagraph_capture_sizes = sorted(set(cudagraph_capture_sizes))
|
||||
|
||||
|
||||
@@ -584,8 +584,6 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
|
||||
top_k: int,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
dispatch_dtype_bytes_per_elem: int = 0,
|
||||
dispatch_scale_bytes_per_token: int = 0,
|
||||
):
|
||||
"""Initialize the MoeAlltoAll workspace."""
|
||||
if self.initialized:
|
||||
@@ -616,13 +614,9 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
|
||||
ep_config = MnnvlConfig(
|
||||
comm_backend=CustomCommunicator(self.cpu_group),
|
||||
)
|
||||
if dispatch_dtype_bytes_per_elem == 0:
|
||||
hidden_bytes = hidden_size // 2
|
||||
else:
|
||||
hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem
|
||||
total_dispatch_payload_size_per_token = (
|
||||
hidden_bytes
|
||||
+ dispatch_scale_bytes_per_token
|
||||
hidden_size // 2 # nvfp4 hidden states
|
||||
+ hidden_size // 16 # fp8 scaling factors
|
||||
+ top_k * 4 # int32 topks ids
|
||||
+ top_k * 4 # float32 topk weights
|
||||
)
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Scheduler-side logic for the NIXL connector."""
|
||||
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any
|
||||
@@ -88,6 +89,12 @@ class NixlConnectorScheduler:
|
||||
if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
|
||||
logger.info("Hybrid Memory Allocator is enabled with NIXL")
|
||||
|
||||
if os.environ.get("VLLM_NIXL_ABORT_REQUEST_TIMEOUT") is not None:
|
||||
logger.warning(
|
||||
"VLLM_NIXL_ABORT_REQUEST_TIMEOUT is deprecated and will be "
|
||||
"removed in release 0.22.0."
|
||||
)
|
||||
|
||||
# Background thread for handling new handshake requests.
|
||||
self._nixl_handshake_listener_t: threading.Thread | None = None
|
||||
self._stop_event = threading.Event()
|
||||
|
||||
@@ -542,6 +542,14 @@ class EngineArgs:
|
||||
mm_encoder_attn_backend: AttentionBackendEnum | str | None = (
|
||||
MultiModalConfig.mm_encoder_attn_backend
|
||||
)
|
||||
mm_encoder_attn_dtype: str | None = MultiModalConfig.mm_encoder_attn_dtype
|
||||
mm_encoder_fp8_scale_path: str | None = MultiModalConfig.mm_encoder_fp8_scale_path
|
||||
mm_encoder_fp8_scale_save_path: str | None = (
|
||||
MultiModalConfig.mm_encoder_fp8_scale_save_path
|
||||
)
|
||||
mm_encoder_fp8_scale_save_margin: float = (
|
||||
MultiModalConfig.mm_encoder_fp8_scale_save_margin
|
||||
)
|
||||
io_processor_plugin: str | None = None
|
||||
renderer_num_workers: int = 1
|
||||
skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling
|
||||
@@ -1179,6 +1187,22 @@ class EngineArgs:
|
||||
"--mm-encoder-attn-backend",
|
||||
**multimodal_kwargs["mm_encoder_attn_backend"],
|
||||
)
|
||||
multimodal_group.add_argument(
|
||||
"--mm-encoder-attn-dtype",
|
||||
**multimodal_kwargs["mm_encoder_attn_dtype"],
|
||||
)
|
||||
multimodal_group.add_argument(
|
||||
"--mm-encoder-fp8-scale-path",
|
||||
**multimodal_kwargs["mm_encoder_fp8_scale_path"],
|
||||
)
|
||||
multimodal_group.add_argument(
|
||||
"--mm-encoder-fp8-scale-save-path",
|
||||
**multimodal_kwargs["mm_encoder_fp8_scale_save_path"],
|
||||
)
|
||||
multimodal_group.add_argument(
|
||||
"--mm-encoder-fp8-scale-save-margin",
|
||||
**multimodal_kwargs["mm_encoder_fp8_scale_save_margin"],
|
||||
)
|
||||
multimodal_group.add_argument(
|
||||
"--interleave-mm-strings", **multimodal_kwargs["interleave_mm_strings"]
|
||||
)
|
||||
@@ -1517,6 +1541,10 @@ class EngineArgs:
|
||||
mm_encoder_only=self.mm_encoder_only,
|
||||
mm_encoder_tp_mode=self.mm_encoder_tp_mode,
|
||||
mm_encoder_attn_backend=self.mm_encoder_attn_backend,
|
||||
mm_encoder_attn_dtype=self.mm_encoder_attn_dtype,
|
||||
mm_encoder_fp8_scale_path=self.mm_encoder_fp8_scale_path,
|
||||
mm_encoder_fp8_scale_save_path=self.mm_encoder_fp8_scale_save_path,
|
||||
mm_encoder_fp8_scale_save_margin=self.mm_encoder_fp8_scale_save_margin,
|
||||
pooler_config=self.pooler_config,
|
||||
generation_config=self.generation_config,
|
||||
override_generation_config=self.override_generation_config,
|
||||
|
||||
@@ -78,7 +78,6 @@ class EngineClient(ABC):
|
||||
priority: int = 0,
|
||||
data_parallel_rank: int | None = None,
|
||||
reasoning_ended: bool | None = None,
|
||||
reasoning_parser_kwargs: dict[str, Any] | None = None,
|
||||
) -> AsyncGenerator[RequestOutput, None]:
|
||||
"""Generate outputs for a request."""
|
||||
...
|
||||
|
||||
@@ -79,7 +79,7 @@ from vllm.renderers.inputs.preprocess import (
|
||||
prompt_to_seq,
|
||||
)
|
||||
from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
|
||||
from vllm.tasks import PoolingTask
|
||||
from vllm.tasks import SCORE_TYPE_MAP, PoolingTask
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.usage.usage_lib import UsageContext
|
||||
from vllm.utils.counter import Counter
|
||||
@@ -1204,12 +1204,9 @@ class LLM:
|
||||
f"Supported tasks: {self.supported_tasks}"
|
||||
)
|
||||
else:
|
||||
logger.warning_once(
|
||||
"Pooling multitask support is deprecated and will "
|
||||
"be removed in v0.20. When the default pooling task is "
|
||||
"not what you want, you need to manually specify it "
|
||||
'via PoolerConfig(task="%s"). ',
|
||||
pooling_task,
|
||||
raise ValueError(
|
||||
f"Try switching the model's pooling_task "
|
||||
f'via `PoolerConfig(task="{pooling_task}")`'
|
||||
)
|
||||
|
||||
if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors:
|
||||
@@ -1412,7 +1409,7 @@ class LLM:
|
||||
"pooling model."
|
||||
)
|
||||
|
||||
score_type = self.model_config.score_type
|
||||
score_type: str | None = SCORE_TYPE_MAP.get(self.pooling_task, None) # type: ignore[arg-type]
|
||||
if (
|
||||
score_type == "cross-encoder"
|
||||
and getattr(self.model_config.hf_config, "num_labels", 0) != 1
|
||||
|
||||
@@ -317,4 +317,5 @@ class OpenAIServingChatBatch(OpenAIServingChat):
|
||||
model=model_name,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
system_fingerprint=self.system_fingerprint,
|
||||
)
|
||||
|
||||
@@ -129,6 +129,9 @@ class ChatCompletionStreamResponse(OpenAIBaseModel):
|
||||
model: str
|
||||
choices: list[ChatCompletionResponseStreamChoice]
|
||||
usage: UsageInfo | None = Field(default=None)
|
||||
# Set only on the final chunk of a stream to mirror non-streaming responses
|
||||
# without the per-chunk serialization overhead.
|
||||
system_fingerprint: str | None = None
|
||||
# not part of the OpenAI spec but for tracing the tokens
|
||||
prompt_token_ids: list[int] | None = None
|
||||
|
||||
|
||||
@@ -347,11 +347,6 @@ class OpenAIServingChat(OpenAIServing):
|
||||
priority=request.priority,
|
||||
data_parallel_rank=data_parallel_rank,
|
||||
reasoning_ended=reasoning_ended,
|
||||
reasoning_parser_kwargs={
|
||||
"chat_template_kwargs": chat_template_kwargs,
|
||||
}
|
||||
if reasoning_parser
|
||||
else None,
|
||||
)
|
||||
|
||||
generators.append(generator)
|
||||
@@ -1200,6 +1195,16 @@ class OpenAIServingChat(OpenAIServing):
|
||||
choices=[choice_data],
|
||||
model=model_name,
|
||||
)
|
||||
# Stamp the fingerprint on terminal chunks only (those with
|
||||
# finish_reason set). When ``include_usage`` is on, the
|
||||
# trailing usage chunk below overrides this as the true
|
||||
# final message.
|
||||
if (
|
||||
not include_usage
|
||||
and self.system_fingerprint is not None
|
||||
and choice_data.finish_reason is not None
|
||||
):
|
||||
chunk.system_fingerprint = self.system_fingerprint
|
||||
|
||||
# handle usage stats if requested & if continuous
|
||||
if include_continuous_usage:
|
||||
@@ -1234,6 +1239,7 @@ class OpenAIServingChat(OpenAIServing):
|
||||
choices=[],
|
||||
model=model_name,
|
||||
usage=final_usage,
|
||||
system_fingerprint=self.system_fingerprint,
|
||||
)
|
||||
final_usage_data = final_usage_chunk.model_dump_json(
|
||||
exclude_unset=True, exclude_none=True
|
||||
@@ -1642,6 +1648,7 @@ class OpenAIServingChat(OpenAIServing):
|
||||
model=model_name,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
system_fingerprint=self.system_fingerprint,
|
||||
prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
|
||||
prompt_token_ids=(
|
||||
final_res.prompt_token_ids if request.return_token_ids else None
|
||||
|
||||
@@ -153,9 +153,21 @@ class BaseFrontendArgs:
|
||||
"""If set to True, log the stack trace of error responses"""
|
||||
tokens_only: bool = False
|
||||
"""
|
||||
If set to True, only enable the Tokens In<>Out endpoint.
|
||||
If set to True, only enable the Tokens In<>Out endpoint.
|
||||
This is intended for use in a Disaggregated Everything setup.
|
||||
"""
|
||||
fingerprint_mode: Literal["full", "hash", "custom", "none"] = "full"
|
||||
"""Controls the ``system_fingerprint`` field on responses.
|
||||
|
||||
- ``full`` (default): ``vllm-<version>[-<parallelism>]-<hash8>``. Encodes
|
||||
server version, non-trivial parallelism degrees (tp/pp/dp/ep), and an
|
||||
8-char config hash.
|
||||
- ``hash``: ``vllm-<version>-<hash8>``. Parallelism stripped.
|
||||
- ``custom``: emits the literal string from ``--fingerprint-value``.
|
||||
- ``none``: the field is omitted (serialized as ``null``).
|
||||
"""
|
||||
fingerprint_value: str | None = None
|
||||
"""Literal fingerprint string used when ``--fingerprint-mode=custom``."""
|
||||
|
||||
@classmethod
|
||||
def _customize_cli_kwargs(
|
||||
|
||||
@@ -512,3 +512,6 @@ class CompletionStreamResponse(OpenAIBaseModel):
|
||||
model: str
|
||||
choices: list[CompletionResponseStreamChoice]
|
||||
usage: UsageInfo | None = Field(default=None)
|
||||
# Set only on the final chunk of a stream to mirror non-streaming responses
|
||||
# without the per-chunk serialization overhead.
|
||||
system_fingerprint: str | None = None
|
||||
|
||||
@@ -383,6 +383,7 @@ class OpenAIServingCompletion(OpenAIServing):
|
||||
|
||||
chunk = CompletionStreamResponse(
|
||||
id=request_id,
|
||||
object="text_completion",
|
||||
created=created_time,
|
||||
model=model_name,
|
||||
choices=[
|
||||
@@ -401,6 +402,14 @@ class OpenAIServingCompletion(OpenAIServing):
|
||||
)
|
||||
],
|
||||
)
|
||||
# Stamp on terminal chunk only when no trailing usage chunk
|
||||
# will follow (that one is the true final message).
|
||||
if (
|
||||
not include_usage
|
||||
and self.system_fingerprint is not None
|
||||
and finish_reason is not None
|
||||
):
|
||||
chunk.system_fingerprint = self.system_fingerprint
|
||||
if include_continuous_usage:
|
||||
prompt_tokens = num_prompt_tokens[prompt_idx]
|
||||
completion_tokens = previous_num_tokens[i]
|
||||
@@ -410,7 +419,7 @@ class OpenAIServingCompletion(OpenAIServing):
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
)
|
||||
|
||||
response_json = chunk.model_dump_json(exclude_unset=False)
|
||||
response_json = chunk.model_dump_json(exclude_unset=True)
|
||||
yield f"data: {response_json}\n\n"
|
||||
|
||||
total_prompt_tokens = sum(num_prompt_tokens)
|
||||
@@ -433,6 +442,7 @@ class OpenAIServingCompletion(OpenAIServing):
|
||||
model=model_name,
|
||||
choices=[],
|
||||
usage=final_usage_info,
|
||||
system_fingerprint=self.system_fingerprint,
|
||||
)
|
||||
final_usage_data = final_usage_chunk.model_dump_json(
|
||||
exclude_unset=False, exclude_none=True
|
||||
@@ -562,6 +572,7 @@ class OpenAIServingCompletion(OpenAIServing):
|
||||
model=model_name,
|
||||
choices=choices,
|
||||
usage=usage,
|
||||
system_fingerprint=self.system_fingerprint,
|
||||
kv_transfer_params=kv_transfer_params,
|
||||
)
|
||||
|
||||
|
||||
@@ -157,6 +157,19 @@ class OpenAIServing:
|
||||
self.renderer = engine_client.renderer
|
||||
self.input_processor = engine_client.input_processor
|
||||
|
||||
# Computed once at startup (cached by ``vllm_config`` identity) and
|
||||
# stamped on non-streaming responses. Streaming chunks deliberately
|
||||
# omit it to avoid per-chunk overhead.
|
||||
from vllm.entrypoints.openai.fingerprint import get_system_fingerprint
|
||||
|
||||
try:
|
||||
self.system_fingerprint: str | None = get_system_fingerprint(
|
||||
engine_client.vllm_config
|
||||
)
|
||||
except Exception:
|
||||
# Never fail server startup over the fingerprint.
|
||||
self.system_fingerprint = None
|
||||
|
||||
async def beam_search(
|
||||
self,
|
||||
prompt: EngineInput,
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Build the ``system_fingerprint`` string returned by the OpenAI-compatible
|
||||
server.
|
||||
|
||||
Four modes, configured via ``--fingerprint-mode``:
|
||||
|
||||
* ``full`` (default): ``vllm-<version>[-<parallelism>]-<hash8>`` — encodes
|
||||
server version, any non-trivial parallelism degree (tp/pp/dp/ep), and an
|
||||
8-char prefix of ``vllm_config.compute_hash()`` (covers model identity,
|
||||
quant config, speculative, attention backend, etc.).
|
||||
* ``hash``: ``vllm-<version>-<hash8>`` — parallelism stripped.
|
||||
* ``custom``: user-provided literal via ``--fingerprint-value``.
|
||||
* ``none``: the field is omitted (serialized as ``null``).
|
||||
|
||||
``get_system_fingerprint`` is only called at serving-class init (a handful
|
||||
of times per server); each subclass caches the returned string on
|
||||
``self.system_fingerprint``, so per-request cost is one attribute read.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
FingerprintMode = Literal["full", "hash", "custom", "none"]
|
||||
|
||||
_DEFAULT_MODE: FingerprintMode = "full"
|
||||
_CUSTOM_VALUE: str | None = None
|
||||
|
||||
|
||||
def set_default_fingerprint_mode(
|
||||
mode: FingerprintMode,
|
||||
custom_value: str | None = None,
|
||||
) -> None:
|
||||
"""Configure the fingerprint mode for subsequent ``get_system_fingerprint``
|
||||
calls. Called once at server startup."""
|
||||
global _DEFAULT_MODE, _CUSTOM_VALUE
|
||||
_DEFAULT_MODE = mode
|
||||
_CUSTOM_VALUE = custom_value
|
||||
|
||||
|
||||
def get_system_fingerprint(vllm_config: Any) -> str | None:
|
||||
"""Return the fingerprint for ``vllm_config`` using the mode configured by
|
||||
``set_default_fingerprint_mode``."""
|
||||
return build_system_fingerprint(vllm_config, _DEFAULT_MODE, _CUSTOM_VALUE)
|
||||
|
||||
|
||||
def build_system_fingerprint(
|
||||
vllm_config: Any,
|
||||
mode: FingerprintMode = "full",
|
||||
custom_value: str | None = None,
|
||||
) -> str | None:
|
||||
if mode == "none":
|
||||
return None
|
||||
if mode == "custom":
|
||||
return custom_value
|
||||
|
||||
from vllm import __version__ as vllm_version
|
||||
|
||||
try:
|
||||
hash8 = vllm_config.compute_hash()[:8]
|
||||
except Exception:
|
||||
hash8 = "nohash"
|
||||
|
||||
if mode == "hash":
|
||||
return f"vllm-{vllm_version}-{hash8}"
|
||||
|
||||
# mode == "full"
|
||||
parts: list[str] = [f"vllm-{vllm_version}"]
|
||||
pc = getattr(vllm_config, "parallel_config", None)
|
||||
if pc is not None:
|
||||
tp = getattr(pc, "tensor_parallel_size", 1)
|
||||
if tp > 1:
|
||||
parts.append(f"tp{tp}")
|
||||
pp = getattr(pc, "pipeline_parallel_size", 1)
|
||||
if pp > 1:
|
||||
parts.append(f"pp{pp}")
|
||||
dp = getattr(pc, "data_parallel_size", 1)
|
||||
if dp > 1:
|
||||
parts.append(f"dp{dp}")
|
||||
if getattr(pc, "enable_expert_parallel", False):
|
||||
parts.append("ep")
|
||||
parts.append(hash8)
|
||||
return "-".join(parts)
|
||||
@@ -61,9 +61,17 @@ async def init_generate_state(
|
||||
)
|
||||
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
|
||||
from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
|
||||
from vllm.entrypoints.openai.fingerprint import set_default_fingerprint_mode
|
||||
from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
|
||||
from vllm.entrypoints.serve.disagg.serving import ServingTokens
|
||||
|
||||
# Applied before any serving class is constructed so that each one picks
|
||||
# up the chosen mode on its first cache miss.
|
||||
set_default_fingerprint_mode(
|
||||
getattr(args, "fingerprint_mode", "full"),
|
||||
getattr(args, "fingerprint_value", None),
|
||||
)
|
||||
|
||||
if args.tool_server == "demo":
|
||||
tool_server: ToolServer | None = DemoToolServer()
|
||||
assert isinstance(tool_server, DemoToolServer)
|
||||
|
||||
@@ -472,13 +472,9 @@ class OpenAIServingResponses(OpenAIServing):
|
||||
context = SimpleContext()
|
||||
|
||||
if self.parser and self.parser.reasoning_parser_cls is not None:
|
||||
chat_template_kwargs = self._effective_chat_template_kwargs(request)
|
||||
reasoning_parser_kwargs = {
|
||||
"chat_template_kwargs": chat_template_kwargs,
|
||||
}
|
||||
reasoning_parser = self.parser.reasoning_parser_cls(
|
||||
tokenizer,
|
||||
chat_template_kwargs=chat_template_kwargs,
|
||||
chat_template_kwargs=self._effective_chat_template_kwargs(request),
|
||||
)
|
||||
if (
|
||||
isinstance(
|
||||
@@ -501,9 +497,6 @@ class OpenAIServingResponses(OpenAIServing):
|
||||
lora_request=lora_request,
|
||||
priority=request.priority,
|
||||
trace_headers=trace_headers,
|
||||
reasoning_parser_kwargs=reasoning_parser_kwargs
|
||||
if self.parser and self.parser.reasoning_parser_cls is not None
|
||||
else None,
|
||||
)
|
||||
generators.append(generator)
|
||||
|
||||
@@ -650,7 +643,6 @@ class OpenAIServingResponses(OpenAIServing):
|
||||
lora_request: LoRARequest | None = None,
|
||||
priority: int = 0,
|
||||
trace_headers: Mapping[str, str] | None = None,
|
||||
reasoning_parser_kwargs: dict[str, Any] | None = None,
|
||||
):
|
||||
max_model_len = self.model_config.max_model_len
|
||||
|
||||
@@ -674,7 +666,6 @@ class OpenAIServingResponses(OpenAIServing):
|
||||
lora_request=lora_request,
|
||||
trace_headers=trace_headers,
|
||||
priority=priority,
|
||||
reasoning_parser_kwargs=reasoning_parser_kwargs,
|
||||
)
|
||||
|
||||
async for res in generator:
|
||||
|
||||
@@ -15,10 +15,7 @@ from starlette.datastructures import Headers
|
||||
from vllm import PoolingParams, PoolingRequestOutput, envs
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.engine.protocol import EngineClient
|
||||
from vllm.entrypoints.chat_utils import (
|
||||
ChatTemplateConfig,
|
||||
ChatTemplateContentFormatOption,
|
||||
)
|
||||
from vllm.entrypoints.chat_utils import ChatTemplateConfig
|
||||
from vllm.entrypoints.logger import RequestLogger
|
||||
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
|
||||
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
|
||||
@@ -48,9 +45,7 @@ class PoolingServingBase(ABC):
|
||||
models: OpenAIServingModels,
|
||||
*,
|
||||
request_logger: RequestLogger | None,
|
||||
chat_template: str | None = None,
|
||||
chat_template_content_format: ChatTemplateContentFormatOption = "auto",
|
||||
trust_request_chat_template: bool = False,
|
||||
chat_template_config: ChatTemplateConfig,
|
||||
return_tokens_as_token_ids: bool = False,
|
||||
log_error_stack: bool = False,
|
||||
):
|
||||
@@ -63,11 +58,7 @@ class PoolingServingBase(ABC):
|
||||
self.request_logger = request_logger
|
||||
self.return_tokens_as_token_ids = return_tokens_as_token_ids
|
||||
self.log_error_stack = log_error_stack
|
||||
self.chat_template_config = ChatTemplateConfig(
|
||||
chat_template=chat_template,
|
||||
chat_template_content_format=chat_template_content_format,
|
||||
trust_request_chat_template=trust_request_chat_template,
|
||||
)
|
||||
self.chat_template_config = chat_template_config
|
||||
|
||||
# Shared thread pool executor for preprocessing and postprocessing.
|
||||
self._executor: Executor = models.renderer._executor
|
||||
|
||||
@@ -10,7 +10,7 @@ from vllm.entrypoints.chat_utils import ChatTemplateConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.plugins.io_processors import has_io_processor
|
||||
from vllm.renderers import BaseRenderer
|
||||
from vllm.tasks import POOLING_TASKS, SupportedTask
|
||||
from vllm.tasks import POOLING_TASKS, SCORE_TYPE_MAP, SupportedTask
|
||||
|
||||
from .base.io_processor import PoolingIOProcessor
|
||||
from .utils import enable_scoring_api
|
||||
@@ -43,23 +43,24 @@ def init_pooling_io_processors(
|
||||
) -> dict[str, PoolingIOProcessor]:
|
||||
model_config = vllm_config.model_config
|
||||
processors: dict[str, type[PoolingIOProcessor]] = {}
|
||||
pooling_task = model_config.get_pooling_task(supported_tasks)
|
||||
|
||||
if "classify" in supported_tasks:
|
||||
if pooling_task == "classify":
|
||||
from .classify.io_processor import ClassifyIOProcessor
|
||||
|
||||
processors["classify"] = ClassifyIOProcessor
|
||||
|
||||
if "token_classify" in supported_tasks:
|
||||
if pooling_task == "token_classify":
|
||||
from .classify.io_processor import TokenClassifyIOProcessor
|
||||
|
||||
processors["token_classify"] = TokenClassifyIOProcessor
|
||||
|
||||
if "embed" in supported_tasks:
|
||||
if pooling_task == "embed":
|
||||
from .embed.io_processor import EmbedIOProcessor
|
||||
|
||||
processors["embed"] = EmbedIOProcessor
|
||||
|
||||
if "token_embed" in supported_tasks:
|
||||
if pooling_task == "token_embed":
|
||||
from .embed.io_processor import TokenEmbedIOProcessor
|
||||
|
||||
processors["token_embed"] = TokenEmbedIOProcessor
|
||||
@@ -71,15 +72,15 @@ def init_pooling_io_processors(
|
||||
from .pooling.io_processor import PluginWithIOProcessorPlugins
|
||||
|
||||
processors["plugin"] = PluginWithIOProcessorPlugins
|
||||
elif "plugin" in supported_tasks:
|
||||
elif pooling_task == "plugin":
|
||||
from .pooling.io_processor import PluginWithoutIOProcessorPlugins
|
||||
|
||||
processors["plugin"] = PluginWithoutIOProcessorPlugins
|
||||
|
||||
if enable_scoring_api(supported_tasks, model_config):
|
||||
score_type = model_config.score_type
|
||||
from .scoring.io_processor import ScoringIOProcessors
|
||||
|
||||
score_type: str | None = SCORE_TYPE_MAP.get(pooling_task, None) # type: ignore[arg-type]
|
||||
if score_type is not None and score_type in ScoringIOProcessors:
|
||||
processors[score_type] = ScoringIOProcessors[score_type]
|
||||
|
||||
@@ -140,6 +141,10 @@ def init_pooling_state(
|
||||
request_logger: RequestLogger | None,
|
||||
supported_tasks: tuple["SupportedTask", ...],
|
||||
):
|
||||
model_config = engine_client.model_config
|
||||
if model_config is None:
|
||||
return
|
||||
|
||||
from vllm.entrypoints.chat_utils import load_chat_template
|
||||
from vllm.tasks import POOLING_TASKS
|
||||
|
||||
@@ -148,8 +153,14 @@ def init_pooling_state(
|
||||
from .pooling.serving import ServingPooling
|
||||
from .scoring.serving import ServingScores
|
||||
|
||||
model_config = engine_client.model_config
|
||||
resolved_chat_template = load_chat_template(args.chat_template)
|
||||
pooling_task = model_config.get_pooling_task(supported_tasks)
|
||||
|
||||
chat_template_config = ChatTemplateConfig(
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
)
|
||||
|
||||
state.serving_pooling = (
|
||||
(
|
||||
@@ -158,9 +169,7 @@ def init_pooling_state(
|
||||
state.openai_serving_models,
|
||||
supported_tasks=supported_tasks,
|
||||
request_logger=request_logger,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
chat_template_config=chat_template_config,
|
||||
)
|
||||
)
|
||||
if any(t in supported_tasks for t in POOLING_TASKS)
|
||||
@@ -171,11 +180,9 @@ def init_pooling_state(
|
||||
engine_client,
|
||||
state.openai_serving_models,
|
||||
request_logger=request_logger,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
chat_template_config=chat_template_config,
|
||||
)
|
||||
if "embed" in supported_tasks
|
||||
if pooling_task == "embed"
|
||||
else None
|
||||
)
|
||||
state.serving_classification = (
|
||||
@@ -183,21 +190,18 @@ def init_pooling_state(
|
||||
engine_client,
|
||||
state.openai_serving_models,
|
||||
request_logger=request_logger,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
chat_template_config=chat_template_config,
|
||||
)
|
||||
if "classify" in supported_tasks
|
||||
if pooling_task == "classify"
|
||||
else None
|
||||
)
|
||||
state.serving_scores = (
|
||||
ServingScores(
|
||||
engine_client,
|
||||
state.openai_serving_models,
|
||||
supported_tasks=supported_tasks,
|
||||
request_logger=request_logger,
|
||||
chat_template=resolved_chat_template,
|
||||
chat_template_content_format=args.chat_template_content_format,
|
||||
trust_request_chat_template=args.trust_request_chat_template,
|
||||
chat_template_config=chat_template_config,
|
||||
enable_flash_late_interaction=getattr(
|
||||
args, "enable_flash_late_interaction", True
|
||||
),
|
||||
@@ -214,7 +218,12 @@ def get_pooling_invocation_types(
|
||||
# NOTE: Items defined earlier take higher priority
|
||||
invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = []
|
||||
|
||||
if "embed" in supported_tasks:
|
||||
if model_config is None:
|
||||
return invocation_types
|
||||
|
||||
pooling_task = model_config.get_pooling_task(supported_tasks)
|
||||
|
||||
if pooling_task == "embed":
|
||||
from .embed.api_router import create_embedding, embedding
|
||||
from .embed.protocol import EmbeddingRequest
|
||||
|
||||
@@ -222,7 +231,7 @@ def get_pooling_invocation_types(
|
||||
(EmbeddingRequest, (embedding, create_embedding)),
|
||||
]
|
||||
|
||||
if "classify" in supported_tasks:
|
||||
if pooling_task == "classify":
|
||||
from .classify.api_router import classify, create_classify
|
||||
from .classify.protocol import ClassificationRequest
|
||||
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user