forked from Karylab-cklius/vllm
Compare commits
80
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a26e8dc7ff | ||
|
|
599e7359a3 | ||
|
|
d0cf73ce42 | ||
|
|
f0a5c5973b | ||
|
|
b7e4b88987 | ||
|
|
90b29e5302 | ||
|
|
a45d96ff42 | ||
|
|
7693c8eabf | ||
|
|
7624525bf6 | ||
|
|
d1b4f10b19 | ||
|
|
9fdc0f3aeb | ||
|
|
05d96d7991 | ||
|
|
ccbc5ac449 | ||
|
|
bcf2be9612 | ||
|
|
89138b21cc | ||
|
|
6edd43de3c | ||
|
|
16c971dbc7 | ||
|
|
262ddd0d81 | ||
|
|
e60c1674b3 | ||
|
|
faa80947f5 | ||
|
|
eeabf740bb | ||
|
|
cdcffafef8 | ||
|
|
4d22667c32 | ||
|
|
1fe3932c8b | ||
|
|
2dccb38f73 | ||
|
|
d157216093 | ||
|
|
93f3c8e531 | ||
|
|
2cc26c3a99 | ||
|
|
dfa8852db2 | ||
|
|
714c6e0eab | ||
|
|
0fefd00e6c | ||
|
|
f5c081d432 | ||
|
|
c88ea8338b | ||
|
|
9f9ecff4cd | ||
|
|
ca1954d58c | ||
|
|
55e6d3d5c0 | ||
|
|
6682c231fa | ||
|
|
5ae685c1c8 | ||
|
|
ce8cf9161d | ||
|
|
18be11fd59 | ||
|
|
8d8855fdae | ||
|
|
e855d380fa | ||
|
|
0e5a9382af | ||
|
|
04bf5a35fa | ||
|
|
43a73f853b | ||
|
|
ffbc2e5bdb | ||
|
|
f9e6db3034 | ||
|
|
d61d2b08e9 | ||
|
|
f5e59ee7a6 | ||
|
|
9b005edc48 | ||
|
|
bf9a185395 | ||
|
|
ad041c79db | ||
|
|
747b068136 | ||
|
|
122f75d939 | ||
|
|
d8f8a7aad2 | ||
|
|
0115e957d4 | ||
|
|
116ed130f4 | ||
|
|
8374387bd8 | ||
|
|
912fbe9555 | ||
|
|
52131f88d9 | ||
|
|
821eb80c0d | ||
|
|
a2956a0f8e | ||
|
|
911355e216 | ||
|
|
8d3f8f485e | ||
|
|
96efb91480 | ||
|
|
2754231ba3 | ||
|
|
2390d44209 | ||
|
|
7362b4450a | ||
|
|
57a314d155 | ||
|
|
d4c57863f7 | ||
|
|
68e1b711f1 | ||
|
|
0024f39a32 | ||
|
|
e9163b536e | ||
|
|
7acaea634c | ||
|
|
697e4ff352 | ||
|
|
a3e2e250f0 | ||
|
|
143e4dccdf | ||
|
|
6590a3ecda | ||
|
|
b3debb7e77 | ||
|
|
458c1a4b2d |
@@ -12,7 +12,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cuda-12-9
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
@@ -27,7 +27,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cuda-13-0
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
@@ -42,7 +42,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cpu
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
|
||||
- "mkdir artifacts"
|
||||
@@ -55,7 +55,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cuda-12-9
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
@@ -68,7 +68,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cuda-13-0
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
@@ -81,7 +81,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cpu
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
|
||||
- "mkdir artifacts"
|
||||
@@ -97,7 +97,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-release-image-x86
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
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"
|
||||
- "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 FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
@@ -110,7 +110,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
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"
|
||||
- "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 FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
@@ -120,7 +120,7 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-cuda-13-0
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
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"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
@@ -133,13 +133,57 @@ steps:
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-cuda-13-0
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
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"
|
||||
# compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-ubuntu2404
|
||||
agents:
|
||||
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"
|
||||
- "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 FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --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"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-ubuntu2404
|
||||
agents:
|
||||
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"
|
||||
- "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 FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --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 13.0 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-cuda-13-0-ubuntu2404
|
||||
agents:
|
||||
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"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-cuda-13-0-ubuntu2404
|
||||
agents:
|
||||
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"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
|
||||
|
||||
- block: "Build release image for x86_64 CPU"
|
||||
key: block-cpu-release-image-build
|
||||
depends_on: ~
|
||||
@@ -148,8 +192,9 @@ steps:
|
||||
depends_on:
|
||||
- block-cpu-release-image-build
|
||||
- input-release-version
|
||||
id: build-release-image-x86-cpu
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
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"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
|
||||
@@ -163,11 +208,12 @@ steps:
|
||||
depends_on: ~
|
||||
|
||||
- label: "Build release image - arm64 - CPU"
|
||||
depends_on:
|
||||
depends_on:
|
||||
- block-arm64-cpu-release-image-build
|
||||
- input-release-version
|
||||
id: build-release-image-arm64-cpu
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
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"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version) --tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest --progress plain --target vllm-openai -f docker/Dockerfile.cpu ."
|
||||
@@ -185,7 +231,7 @@ steps:
|
||||
- build-release-image-arm64
|
||||
id: create-multi-arch-manifest
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
|
||||
@@ -196,7 +242,7 @@ steps:
|
||||
- create-multi-arch-manifest
|
||||
id: annotate-release-workflow
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-release.sh"
|
||||
|
||||
@@ -206,18 +252,67 @@ steps:
|
||||
- build-release-image-arm64-cuda-13-0
|
||||
id: create-multi-arch-manifest-cuda-13-0
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on:
|
||||
- build-release-image-x86-ubuntu2404
|
||||
- build-release-image-arm64-ubuntu2404
|
||||
id: create-multi-arch-manifest-ubuntu2404
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on:
|
||||
- build-release-image-x86-cuda-13-0-ubuntu2404
|
||||
- build-release-image-arm64-cuda-13-0-ubuntu2404
|
||||
id: create-multi-arch-manifest-cuda-13-0-ubuntu2404
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
|
||||
- block: "Confirm publishing release images to DockerHub"
|
||||
key: block-publish-release-images-dockerhub
|
||||
depends_on:
|
||||
- create-multi-arch-manifest
|
||||
- create-multi-arch-manifest-cuda-13-0
|
||||
- build-release-image-x86-cpu
|
||||
- build-release-image-arm64-cpu
|
||||
- build-rocm-release-image
|
||||
|
||||
- label: "Publish release images to DockerHub"
|
||||
key: publish-release-images-dockerhub
|
||||
depends_on:
|
||||
- block-publish-release-images-dockerhub
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/push-release-builds.sh"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
password-env: DOCKERHUB_TOKEN
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
|
||||
- label: "Publish nightly multi-arch image to DockerHub"
|
||||
depends_on:
|
||||
- create-multi-arch-manifest
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/push-nightly-builds.sh"
|
||||
# Clean up old nightly builds (keep only last 14)
|
||||
@@ -235,7 +330,7 @@ steps:
|
||||
- create-multi-arch-manifest-cuda-13-0
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/push-nightly-builds.sh cu130"
|
||||
# Clean up old nightly builds (keep only last 14)
|
||||
@@ -262,7 +357,7 @@ steps:
|
||||
- block-upload-release-wheels
|
||||
id: upload-release-wheels
|
||||
agents:
|
||||
queue: small_cpu_queue_postmerge
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
|
||||
|
||||
@@ -323,7 +418,7 @@ steps:
|
||||
- step: input-rocm-config
|
||||
allow_failure: true # Allow failure so non-UI builds can proceed (input step is skipped)
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
# Set configuration and check cache
|
||||
- |
|
||||
@@ -465,7 +560,7 @@ steps:
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
# Download artifacts and prepare Docker image
|
||||
@@ -575,7 +670,7 @@ steps:
|
||||
- step: build-rocm-vllm-wheel
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
# Download all wheel artifacts and run upload
|
||||
@@ -624,7 +719,7 @@ steps:
|
||||
- step: input-release-version
|
||||
allow_failure: true
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-rocm-release.sh"
|
||||
env:
|
||||
@@ -641,7 +736,7 @@ steps:
|
||||
depends_on: block-generate-root-index-rocm-wheels
|
||||
id: generate-root-index-rocm-wheels
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
|
||||
env:
|
||||
@@ -655,7 +750,7 @@ steps:
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- |
|
||||
|
||||
Executable
+113
@@ -0,0 +1,113 @@
|
||||
#!/bin/bash
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
# Ensure git tags are up-to-date (Buildkite's default fetch doesn't always include tags)
|
||||
echo "Fetching latest tags from origin..."
|
||||
git fetch --tags --force origin
|
||||
|
||||
# Derive release version from the git tag on the current commit.
|
||||
# The pipeline must be triggered on a tagged commit (e.g. v0.18.1).
|
||||
RELEASE_VERSION=$(git describe --exact-match --tags "${BUILDKITE_COMMIT}" 2>/dev/null || true)
|
||||
if [ -z "${RELEASE_VERSION}" ]; then
|
||||
echo "[FATAL] Commit ${BUILDKITE_COMMIT} has no exact git tag. " \
|
||||
"Release images must be published from a tagged commit."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Strip leading 'v' for use in Docker tags (e.g. v0.18.1 -> 0.18.1)
|
||||
PURE_VERSION="${RELEASE_VERSION#v}"
|
||||
|
||||
echo "========================================"
|
||||
echo "Publishing release images"
|
||||
echo " Commit: ${BUILDKITE_COMMIT}"
|
||||
echo " Release version: ${RELEASE_VERSION}"
|
||||
echo "========================================"
|
||||
|
||||
set -x
|
||||
|
||||
# ---- CUDA (default, CUDA 12.9) ----
|
||||
docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64"
|
||||
docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64"
|
||||
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64" "vllm/vllm-openai:latest-x86_64"
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64" "vllm/vllm-openai:v${PURE_VERSION}-x86_64"
|
||||
docker push "vllm/vllm-openai:latest-x86_64"
|
||||
docker push "vllm/vllm-openai:v${PURE_VERSION}-x86_64"
|
||||
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64" "vllm/vllm-openai:latest-aarch64"
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64" "vllm/vllm-openai:v${PURE_VERSION}-aarch64"
|
||||
docker push "vllm/vllm-openai:latest-aarch64"
|
||||
docker push "vllm/vllm-openai:v${PURE_VERSION}-aarch64"
|
||||
|
||||
docker manifest rm "vllm/vllm-openai:latest" || true
|
||||
docker manifest create "vllm/vllm-openai:latest" "vllm/vllm-openai:latest-x86_64" "vllm/vllm-openai:latest-aarch64"
|
||||
docker manifest push "vllm/vllm-openai:latest"
|
||||
|
||||
docker manifest rm "vllm/vllm-openai:v${PURE_VERSION}" || true
|
||||
docker manifest create "vllm/vllm-openai:v${PURE_VERSION}" "vllm/vllm-openai:v${PURE_VERSION}-x86_64" "vllm/vllm-openai:v${PURE_VERSION}-aarch64"
|
||||
docker manifest push "vllm/vllm-openai:v${PURE_VERSION}"
|
||||
|
||||
# ---- CUDA 13.0 ----
|
||||
docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130"
|
||||
docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130"
|
||||
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130" "vllm/vllm-openai:latest-x86_64-cu130"
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130" "vllm/vllm-openai:v${PURE_VERSION}-x86_64-cu130"
|
||||
docker push "vllm/vllm-openai:latest-x86_64-cu130"
|
||||
docker push "vllm/vllm-openai:v${PURE_VERSION}-x86_64-cu130"
|
||||
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130" "vllm/vllm-openai:latest-aarch64-cu130"
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130" "vllm/vllm-openai:v${PURE_VERSION}-aarch64-cu130"
|
||||
docker push "vllm/vllm-openai:latest-aarch64-cu130"
|
||||
docker push "vllm/vllm-openai:v${PURE_VERSION}-aarch64-cu130"
|
||||
|
||||
docker manifest rm "vllm/vllm-openai:latest-cu130" || true
|
||||
docker manifest create "vllm/vllm-openai:latest-cu130" "vllm/vllm-openai:latest-x86_64-cu130" "vllm/vllm-openai:latest-aarch64-cu130"
|
||||
docker manifest push "vllm/vllm-openai:latest-cu130"
|
||||
|
||||
docker manifest rm "vllm/vllm-openai:v${PURE_VERSION}-cu130" || true
|
||||
docker manifest create "vllm/vllm-openai:v${PURE_VERSION}-cu130" "vllm/vllm-openai:v${PURE_VERSION}-x86_64-cu130" "vllm/vllm-openai:v${PURE_VERSION}-aarch64-cu130"
|
||||
docker manifest push "vllm/vllm-openai:v${PURE_VERSION}-cu130"
|
||||
|
||||
# ---- ROCm ----
|
||||
docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm"
|
||||
docker pull "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base"
|
||||
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm" "vllm/vllm-openai-rocm:latest"
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm" "vllm/vllm-openai-rocm:v${PURE_VERSION}"
|
||||
docker push "vllm/vllm-openai-rocm:latest"
|
||||
docker push "vllm/vllm-openai-rocm:v${PURE_VERSION}"
|
||||
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base" "vllm/vllm-openai-rocm:latest-base"
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm-base" "vllm/vllm-openai-rocm:v${PURE_VERSION}-base"
|
||||
docker push "vllm/vllm-openai-rocm:latest-base"
|
||||
docker push "vllm/vllm-openai-rocm:v${PURE_VERSION}-base"
|
||||
|
||||
# ---- CPU ----
|
||||
# CPU images in ECR are tagged with the full version including 'v' (e.g. v0.18.1),
|
||||
# matching the value from the Buildkite release-version metadata input.
|
||||
docker pull "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:${RELEASE_VERSION}"
|
||||
docker pull "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:${RELEASE_VERSION}"
|
||||
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:${RELEASE_VERSION}" "vllm/vllm-openai-cpu:latest-x86_64"
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:${RELEASE_VERSION}" "vllm/vllm-openai-cpu:v${PURE_VERSION}-x86_64"
|
||||
docker push "vllm/vllm-openai-cpu:latest-x86_64"
|
||||
docker push "vllm/vllm-openai-cpu:v${PURE_VERSION}-x86_64"
|
||||
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:${RELEASE_VERSION}" "vllm/vllm-openai-cpu:latest-arm64"
|
||||
docker tag "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:${RELEASE_VERSION}" "vllm/vllm-openai-cpu:v${PURE_VERSION}-arm64"
|
||||
docker push "vllm/vllm-openai-cpu:latest-arm64"
|
||||
docker push "vllm/vllm-openai-cpu:v${PURE_VERSION}-arm64"
|
||||
|
||||
docker manifest rm "vllm/vllm-openai-cpu:latest" || true
|
||||
docker manifest create "vllm/vllm-openai-cpu:latest" "vllm/vllm-openai-cpu:latest-x86_64" "vllm/vllm-openai-cpu:latest-arm64"
|
||||
docker manifest push "vllm/vllm-openai-cpu:latest"
|
||||
|
||||
docker manifest rm "vllm/vllm-openai-cpu:v${PURE_VERSION}" || true
|
||||
docker manifest create "vllm/vllm-openai-cpu:v${PURE_VERSION}" "vllm/vllm-openai-cpu:v${PURE_VERSION}-x86_64" "vllm/vllm-openai-cpu:v${PURE_VERSION}-arm64"
|
||||
docker manifest push "vllm/vllm-openai-cpu:v${PURE_VERSION}"
|
||||
|
||||
echo "========================================"
|
||||
echo "Successfully published release images for ${RELEASE_VERSION}"
|
||||
echo "========================================"
|
||||
@@ -45,6 +45,22 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
|
||||
|
||||
- label: LM Eval Qwen3.5 Models (B200)
|
||||
timeout_in_minutes: 120
|
||||
device: b200
|
||||
optional: true
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/models/qwen3_5.py
|
||||
- vllm/model_executor/models/qwen3_5_mtp.py
|
||||
- vllm/transformers_utils/configs/qwen3_5.py
|
||||
- vllm/transformers_utils/configs/qwen3_5_moe.py
|
||||
- vllm/model_executor/models/qwen3_next.py
|
||||
- vllm/model_executor/models/qwen3_next_mtp.py
|
||||
- vllm/model_executor/layers/fla/ops/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
|
||||
|
||||
- label: LM Eval Large Models (H200)
|
||||
timeout_in_minutes: 60
|
||||
device: h200
|
||||
|
||||
+1
-1
@@ -27,7 +27,7 @@ pull_request_rules:
|
||||
Hi @{{author}}, the pre-commit checks have failed. Please run:
|
||||
|
||||
```bash
|
||||
uv pip install pre-commit
|
||||
uv pip install pre-commit>=4.5.1
|
||||
pre-commit install
|
||||
pre-commit run --all-files
|
||||
```
|
||||
|
||||
@@ -418,8 +418,8 @@ def _run_single_benchmark(
|
||||
mem_stats = {}
|
||||
if config.profile_memory:
|
||||
mem_stats = {
|
||||
"allocated_mb": torch.cuda.memory_allocated(device) / 1024**2,
|
||||
"reserved_mb": torch.cuda.memory_reserved(device) / 1024**2,
|
||||
"allocated_mb": torch.accelerator.memory_allocated(device) / 1024**2,
|
||||
"reserved_mb": torch.accelerator.memory_reserved(device) / 1024**2,
|
||||
}
|
||||
|
||||
return times, mem_stats
|
||||
|
||||
@@ -95,13 +95,16 @@ def create_logits(
|
||||
def measure_memory() -> tuple[int, int]:
|
||||
"""Return (allocated, reserved) memory in bytes."""
|
||||
torch.accelerator.synchronize()
|
||||
return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()
|
||||
return (
|
||||
torch.accelerator.memory_allocated(),
|
||||
torch.accelerator.max_memory_allocated(),
|
||||
)
|
||||
|
||||
|
||||
def reset_memory_stats():
|
||||
"""Reset peak memory statistics."""
|
||||
reset_buffer_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
torch.accelerator.reset_peak_memory_stats()
|
||||
torch.accelerator.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
@@ -196,6 +196,7 @@ __forceinline__ __device__ u32x8_t ld256_cs(const u32x8_t* addr) {
|
||||
return val;
|
||||
#else
|
||||
assert(false && "ld256_cs requires SM100+ with CUDA 12.9+");
|
||||
return u32x8_t{};
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
@@ -109,16 +109,18 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
|
||||
|
||||
#ifndef USE_ROCM
|
||||
int flag = 0;
|
||||
CUDA_CHECK(cuDeviceGetAttribute(
|
||||
CUresult rdma_result = cuDeviceGetAttribute(
|
||||
&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED,
|
||||
device));
|
||||
if (flag) { // support GPUDirect RDMA if possible
|
||||
device);
|
||||
if (rdma_result == CUDA_SUCCESS &&
|
||||
flag) { // support GPUDirect RDMA if possible
|
||||
prop.allocFlags.gpuDirectRDMACapable = 1;
|
||||
}
|
||||
int fab_flag = 0;
|
||||
CUDA_CHECK(cuDeviceGetAttribute(
|
||||
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device));
|
||||
if (fab_flag) { // support fabric handle if possible
|
||||
CUresult fab_result = cuDeviceGetAttribute(
|
||||
&fab_flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, device);
|
||||
if (fab_result == CUDA_SUCCESS &&
|
||||
fab_flag) { // support fabric handle if possible
|
||||
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -73,10 +73,9 @@ void moe_permute(
|
||||
MOE_DISPATCH(input.scalar_type(), [&] {
|
||||
expandInputRowsKernelLauncher<scalar_t>(
|
||||
get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input),
|
||||
get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx),
|
||||
get_ptr<int>(inv_permuted_idx), get_ptr<int>(permuted_idx),
|
||||
get_ptr<int64_t>(expert_first_token_offset), n_token, valid_num_ptr,
|
||||
n_hidden, topk, n_local_expert, stream);
|
||||
get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx),
|
||||
get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset),
|
||||
n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream);
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@@ -57,7 +57,7 @@ void sortAndScanExpert(const int* expert_for_source_row, const int* source_rows,
|
||||
|
||||
template <typename T>
|
||||
void expandInputRowsKernelLauncher(
|
||||
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
|
||||
T const* unpermuted_input, T* permuted_output,
|
||||
int const* expanded_dest_row_to_expanded_source_row,
|
||||
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
|
||||
int64_t const* expert_first_token_offset, int64_t const num_rows,
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
template <typename T, bool CHECK_SKIPPED>
|
||||
__global__ void expandInputRowsKernel(
|
||||
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
|
||||
T const* unpermuted_input, T* permuted_output,
|
||||
int const* expanded_dest_row_to_expanded_source_row,
|
||||
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
|
||||
int64_t const* expert_first_token_offset, int64_t const num_rows,
|
||||
@@ -16,7 +16,6 @@ __global__ void expandInputRowsKernel(
|
||||
int64_t expanded_dest_row = blockIdx.x;
|
||||
int64_t const expanded_source_row =
|
||||
expanded_dest_row_to_expanded_source_row[expanded_dest_row];
|
||||
int expert_id = sorted_experts[expanded_dest_row];
|
||||
|
||||
if (threadIdx.x == 0) {
|
||||
assert(expanded_dest_row <= INT32_MAX);
|
||||
@@ -54,7 +53,7 @@ __global__ void expandInputRowsKernel(
|
||||
|
||||
template <typename T>
|
||||
void expandInputRowsKernelLauncher(
|
||||
T const* unpermuted_input, T* permuted_output, int* sorted_experts,
|
||||
T const* unpermuted_input, T* permuted_output,
|
||||
int const* expanded_dest_row_to_expanded_source_row,
|
||||
int* expanded_source_row_to_expanded_dest_row, int* permuted_idx,
|
||||
int64_t const* expert_first_token_offset, int64_t const num_rows,
|
||||
@@ -70,12 +69,12 @@ void expandInputRowsKernelLauncher(
|
||||
bool is_check_skip = num_valid_tokens_ptr != nullptr;
|
||||
auto func = func_map[is_check_skip];
|
||||
|
||||
func<<<blocks, threads, 0, stream>>>(
|
||||
unpermuted_input, permuted_output, sorted_experts,
|
||||
expanded_dest_row_to_expanded_source_row,
|
||||
expanded_source_row_to_expanded_dest_row, permuted_idx,
|
||||
expert_first_token_offset, num_rows, num_valid_tokens_ptr, cols, k,
|
||||
num_local_experts);
|
||||
func<<<blocks, threads, 0, stream>>>(unpermuted_input, permuted_output,
|
||||
expanded_dest_row_to_expanded_source_row,
|
||||
expanded_source_row_to_expanded_dest_row,
|
||||
permuted_idx, expert_first_token_offset,
|
||||
num_rows, num_valid_tokens_ptr, cols, k,
|
||||
num_local_experts);
|
||||
}
|
||||
|
||||
template <class T, class U>
|
||||
|
||||
+8
-4
@@ -295,10 +295,14 @@ void cutlass_scaled_sparse_mm(torch::Tensor& out, torch::Tensor const& a,
|
||||
|
||||
std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
|
||||
|
||||
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||
torch::Tensor& output_scale,
|
||||
torch::Tensor const& input_scale,
|
||||
bool is_sf_swizzled_layout);
|
||||
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
|
||||
torch::Tensor const& input, torch::Tensor const& input_scale,
|
||||
bool is_sf_swizzled_layout);
|
||||
|
||||
void scaled_fp4_quant_out(torch::Tensor const& input,
|
||||
torch::Tensor const& input_scale,
|
||||
bool is_sf_swizzled_layout, torch::Tensor& output,
|
||||
torch::Tensor& output_scale);
|
||||
|
||||
void scaled_fp4_experts_quant(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
|
||||
@@ -16,6 +16,8 @@
|
||||
|
||||
#include <torch/all.h>
|
||||
|
||||
#include "nvfp4_utils.cuh"
|
||||
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
||||
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
|
||||
@@ -51,9 +53,10 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
|
||||
torch::Tensor const& output_scale_offset_by_experts);
|
||||
#endif
|
||||
|
||||
void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||
torch::Tensor& output_sf, torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout) {
|
||||
void scaled_fp4_quant_out(torch::Tensor const& input,
|
||||
torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout, torch::Tensor& output,
|
||||
torch::Tensor& output_sf) {
|
||||
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
|
||||
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
|
||||
return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
|
||||
@@ -62,6 +65,34 @@ void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
|
||||
torch::Tensor const& input, torch::Tensor const& input_sf,
|
||||
bool is_sf_swizzled_layout) {
|
||||
int64_t n = input.size(-1);
|
||||
int64_t m = input.numel() / n;
|
||||
auto device = input.device();
|
||||
|
||||
// Two fp4 values packed into a uint8
|
||||
auto output = torch::empty(
|
||||
{m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
|
||||
|
||||
torch::Tensor output_sf;
|
||||
if (is_sf_swizzled_layout) {
|
||||
auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
|
||||
output_sf = torch::empty(
|
||||
{sf_m, sf_n},
|
||||
torch::TensorOptions().device(device).dtype(torch::kInt32));
|
||||
} else {
|
||||
output_sf = torch::empty(
|
||||
{m, n / CVT_FP4_SF_VEC_SIZE},
|
||||
torch::TensorOptions().device(device).dtype(torch::kUInt8));
|
||||
}
|
||||
|
||||
scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
|
||||
output_sf);
|
||||
return {output, output_sf};
|
||||
}
|
||||
|
||||
void scaled_fp4_experts_quant(
|
||||
torch::Tensor& output, torch::Tensor& output_scale,
|
||||
torch::Tensor const& input, torch::Tensor const& input_global_scale,
|
||||
|
||||
@@ -18,6 +18,7 @@
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <utility>
|
||||
|
||||
#include "../../cuda_vec_utils.cuh"
|
||||
|
||||
@@ -54,6 +55,18 @@ inline int computeEffectiveRows(int m) {
|
||||
return round_up(m, ROW_TILE);
|
||||
}
|
||||
|
||||
// Compute the shape of the swizzled SF output tensor.
|
||||
// Returns (rounded_m, rounded_n / 4) where:
|
||||
// rounded_m = round_up(m, 128)
|
||||
// rounded_n = round_up(n / CVT_FP4_SF_VEC_SIZE, 4)
|
||||
inline std::pair<int64_t, int64_t> computeSwizzledSFShape(int64_t m,
|
||||
int64_t n) {
|
||||
int64_t rounded_m = round_up(m, static_cast<int64_t>(128));
|
||||
int64_t scale_n = n / CVT_FP4_SF_VEC_SIZE;
|
||||
int64_t rounded_n = round_up(scale_n, static_cast<int64_t>(4));
|
||||
return {rounded_m, rounded_n / 4};
|
||||
}
|
||||
|
||||
// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
|
||||
inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) {
|
||||
uint32_t val;
|
||||
|
||||
+1
-1
@@ -575,7 +575,7 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(
|
||||
// The range of logits within the row.
|
||||
int rowStart = 0;
|
||||
int seq_len = seqLens[rowIdx / next_n];
|
||||
int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1;
|
||||
int rowEnd = max(0, seq_len - next_n + (rowIdx % next_n) + 1);
|
||||
|
||||
// Local pointers to this block
|
||||
if constexpr (!multipleBlocksPerRow && !mergeBlocks) {
|
||||
|
||||
+15
-4
@@ -564,10 +564,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
|
||||
// Compute NVFP4 block quantized tensor.
|
||||
ops.def(
|
||||
"scaled_fp4_quant(Tensor! output, Tensor input,"
|
||||
" Tensor! output_scale, Tensor input_scale, bool "
|
||||
"is_sf_swizzled_layout) -> ()");
|
||||
ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);
|
||||
"scaled_fp4_quant(Tensor input,"
|
||||
" Tensor input_scale, bool "
|
||||
"is_sf_swizzled_layout) -> (Tensor, Tensor)");
|
||||
ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant_func);
|
||||
|
||||
// Out variant
|
||||
// TODO: Add {at::Tag::out_variant} tag and update all call sites
|
||||
// to use the functional variant once vLLM upgrades PyTorch.
|
||||
// See pytorch/pytorch#176117.
|
||||
ops.def(
|
||||
"scaled_fp4_quant.out(Tensor input,"
|
||||
" Tensor input_scale, bool "
|
||||
"is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) "
|
||||
"-> ()");
|
||||
ops.impl("scaled_fp4_quant.out", torch::kCUDA, &scaled_fp4_quant_out);
|
||||
|
||||
// Compute NVFP4 experts quantization.
|
||||
ops.def(
|
||||
|
||||
+28
-4
@@ -24,6 +24,7 @@
|
||||
|
||||
ARG CUDA_VERSION=12.9.1
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG UBUNTU_VERSION=22.04
|
||||
|
||||
# By parameterizing the base images, we allow third-party to use their own
|
||||
# base images. One use case is hermetic builds with base images stored in
|
||||
@@ -38,7 +39,7 @@ ARG PYTHON_VERSION=3.12
|
||||
# version are not backwards compatible with OSes that use an earlier version.
|
||||
ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04
|
||||
# Using cuda base image with minimal dependencies necessary for JIT compilation (FlashInfer, DeepGEMM, EP kernels)
|
||||
ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-base-ubuntu22.04
|
||||
ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-base-ubuntu${UBUNTU_VERSION}
|
||||
|
||||
# By parameterizing the Deadsnakes repository URL, we allow third-party to use
|
||||
# their own mirror. When doing so, we don't benefit from the transparent
|
||||
@@ -111,6 +112,10 @@ RUN apt-get update -y \
|
||||
gcc-10 \
|
||||
g++-10 \
|
||||
&& update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 110 --slave /usr/bin/g++ g++ /usr/bin/g++-10 \
|
||||
# Install python dev headers if available (needed for cmake FindPython on Ubuntu 24.04
|
||||
# which ships cmake 3.28 and requires Development.SABIModule; silently skipped on
|
||||
# Ubuntu 20.04/22.04 where python3.x-dev is not available without a PPA)
|
||||
&& (apt-get install -y --no-install-recommends python${PYTHON_VERSION}-dev 2>/dev/null || true) \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& curl -LsSf https://astral.sh/uv/install.sh | sh \
|
||||
&& $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
|
||||
@@ -507,7 +512,6 @@ RUN apt-get update -y \
|
||||
software-properties-common \
|
||||
curl \
|
||||
sudo \
|
||||
python3-pip \
|
||||
ffmpeg \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
@@ -535,6 +539,7 @@ RUN apt-get update -y \
|
||||
&& 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
|
||||
|
||||
@@ -593,6 +598,25 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
|
||||
&& flashinfer show-config
|
||||
|
||||
# Pre-download FlashInfer TRTLLM BMM headers for air-gapped environments.
|
||||
# At runtime, MoE JIT compilation downloads these from edge.urm.nvidia.com
|
||||
# which fails without internet. This step caches them at build time.
|
||||
RUN python3 <<'PYEOF'
|
||||
from flashinfer.jit import env as jit_env
|
||||
from flashinfer.jit.cubin_loader import download_trtllm_headers, get_cubin
|
||||
from flashinfer.artifacts import ArtifactPath, CheckSumHash
|
||||
|
||||
download_trtllm_headers(
|
||||
'bmm',
|
||||
jit_env.FLASHINFER_CUBIN_DIR / 'flashinfer' / 'trtllm' / 'batched_gemm' / 'trtllmGen_bmm_export',
|
||||
f'{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export',
|
||||
ArtifactPath.TRTLLM_GEN_BMM,
|
||||
get_cubin(f'{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt', CheckSumHash.TRTLLM_GEN_BMM),
|
||||
)
|
||||
|
||||
print('FlashInfer TRTLLM BMM headers downloaded successfully')
|
||||
PYEOF
|
||||
|
||||
# ============================================================
|
||||
# OPENAI API SERVER DEPENDENCIES
|
||||
# Pre-install these to avoid reinstalling on every vLLM wheel rebuild
|
||||
@@ -620,7 +644,7 @@ RUN set -eux; \
|
||||
ARG BITSANDBYTES_VERSION_X86=0.46.1
|
||||
ARG BITSANDBYTES_VERSION_ARM64=0.42.0
|
||||
ARG TIMM_VERSION=">=1.0.17"
|
||||
ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.3"
|
||||
ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.7"
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
|
||||
BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_ARM64}"; \
|
||||
@@ -628,7 +652,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
BITSANDBYTES_VERSION="${BITSANDBYTES_VERSION_X86}"; \
|
||||
fi; \
|
||||
uv pip install --system accelerate hf_transfer modelscope \
|
||||
"bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs]${RUNAI_MODEL_STREAMER_VERSION}"
|
||||
"bitsandbytes>=${BITSANDBYTES_VERSION}" "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"
|
||||
|
||||
# ============================================================
|
||||
# VLLM INSTALLATION (depends on build stage)
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
#
|
||||
# Build targets:
|
||||
# vllm-openai (default): used for serving deployment
|
||||
# vllm-openai-zen: vLLM from source + zentorch from PyPI via vllm[zen]
|
||||
# vllm-test: used for CI tests
|
||||
# vllm-dev: used for development
|
||||
#
|
||||
@@ -222,3 +223,19 @@ LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}"
|
||||
LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
|
||||
######################### ZEN CPU PYPI IMAGE #########################
|
||||
FROM vllm-openai AS vllm-openai-zen
|
||||
|
||||
ARG TARGETARCH
|
||||
|
||||
RUN if [ "$TARGETARCH" != "amd64" ]; then \
|
||||
echo "ERROR: vllm-openai-amd only supports --platform=linux/amd64"; \
|
||||
exit 1; \
|
||||
fi
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install "vllm[zen]"
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
@@ -33,6 +33,10 @@ group "default" {
|
||||
targets = ["openai"]
|
||||
}
|
||||
|
||||
group "all" {
|
||||
targets = ["openai", "openai-ubuntu2404"]
|
||||
}
|
||||
|
||||
# Base targets
|
||||
|
||||
target "_common" {
|
||||
@@ -74,3 +78,29 @@ target "openai" {
|
||||
tags = ["vllm:openai"]
|
||||
output = ["type=docker"]
|
||||
}
|
||||
|
||||
# Ubuntu 24.04 targets
|
||||
|
||||
target "test-ubuntu2404" {
|
||||
inherits = ["_common", "_labels"]
|
||||
target = "test"
|
||||
tags = ["vllm:test-ubuntu24.04"]
|
||||
args = {
|
||||
UBUNTU_VERSION = "24.04"
|
||||
GDRCOPY_OS_VERSION = "Ubuntu24_04"
|
||||
FLASHINFER_AOT_COMPILE = "true"
|
||||
}
|
||||
output = ["type=docker"]
|
||||
}
|
||||
|
||||
target "openai-ubuntu2404" {
|
||||
inherits = ["_common", "_labels"]
|
||||
target = "vllm-openai"
|
||||
tags = ["vllm:openai-ubuntu24.04"]
|
||||
args = {
|
||||
UBUNTU_VERSION = "24.04"
|
||||
GDRCOPY_OS_VERSION = "Ubuntu24_04"
|
||||
FLASHINFER_AOT_COMPILE = "true"
|
||||
}
|
||||
output = ["type=docker"]
|
||||
}
|
||||
|
||||
@@ -7,6 +7,9 @@
|
||||
"PYTHON_VERSION": {
|
||||
"default": "3.12"
|
||||
},
|
||||
"UBUNTU_VERSION": {
|
||||
"default": "22.04"
|
||||
},
|
||||
"BUILD_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:12.9.1-devel-ubuntu20.04"
|
||||
},
|
||||
@@ -83,7 +86,7 @@
|
||||
"default": ">=1.0.17"
|
||||
},
|
||||
"RUNAI_MODEL_STREAMER_VERSION": {
|
||||
"default": ">=0.15.3"
|
||||
"default": ">=0.15.7"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -75,7 +75,7 @@ For an optimized workflow when iterating on C++/CUDA kernels, see the [Increment
|
||||
vLLM uses `pre-commit` to lint and format the codebase. See <https://pre-commit.com/#usage> if `pre-commit` is new to you. Setting up `pre-commit` is as easy as:
|
||||
|
||||
```bash
|
||||
uv pip install pre-commit
|
||||
uv pip install pre-commit>=4.5.1
|
||||
pre-commit install
|
||||
```
|
||||
|
||||
|
||||
@@ -164,18 +164,18 @@ Priority is **1 = highest** (tried first).
|
||||
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
|
||||
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
|
||||
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A |
|
||||
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
|
||||
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
|
||||
| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
|
||||
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A |
|
||||
| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
|
||||
|
||||
> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
|
||||
>
|
||||
@@ -204,14 +204,14 @@ configuration.
|
||||
|
||||
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
|
||||
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
|
||||
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
|
||||
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
|
||||
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
|
||||
|
||||
@@ -35,7 +35,8 @@ th {
|
||||
| naive | standard | all<sup>1</sup> | G,A,T | N | <sup>6</sup> | [layer.py][vllm.model_executor.layers.fused_moe.layer.FusedMoE] |
|
||||
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize.DeepEPHTPrepareAndFinalize] |
|
||||
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize.DeepEPLLPrepareAndFinalize] |
|
||||
| flashinfer_all2allv | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferA2APrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize.FlashInferA2APrepareAndFinalize] |
|
||||
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
|
||||
!!! info "Table key"
|
||||
1. All types: mxfp4, nvfp4, int4, int8, fp8
|
||||
|
||||
@@ -34,9 +34,6 @@ relies on caching artifacts to reduce start time, we must properly propagate the
|
||||
with the LLM text-backbone, or other instances of the same artifact (as is the case with vision block). `is_encoder=True` is also needed for encoder
|
||||
components (see Compile Range Integration).
|
||||
|
||||
3. `with set_forward_context` context manager should be used around the nn.Module's forward call. This will properly forward the vllm_config which is needed
|
||||
for torch.compile integration.
|
||||
|
||||
### CompilationConfig
|
||||
|
||||
With the exception of `compile_mm_encoder: true`, the multimodal encoder will inherit from the same compilation config as the text LLM. We may extend
|
||||
|
||||
@@ -219,7 +219,7 @@ Supported models:
|
||||
|
||||
* `ibm-granite/granite-4.0-h-small` and other Granite 4.0 models
|
||||
|
||||
Recommended flags: `--tool-call-parser hermes`
|
||||
Recommended flags: `--tool-call-parser granite4`
|
||||
|
||||
* `ibm-granite/granite-3.0-8b-instruct`
|
||||
|
||||
|
||||
@@ -16,4 +16,6 @@ vLLM supports the following hardware platforms:
|
||||
|
||||
vLLM supports third-party hardware plugins that live **outside** the main `vllm` repository. These follow the [Hardware-Pluggable RFC](../../design/plugin_system.md).
|
||||
|
||||
A list of all supported hardware can be found on the [vllm.ai website](https://vllm.ai/#compatibility). If you want to add new hardware, please contact us on [Slack](https://slack.vllm.ai/) or [Email](mailto:collaboration@vllm.ai).
|
||||
A list of all supported hardware can be found on the vLLM website, see [Universal Compatibility - Hardware](https://vllm.ai/#compatibility).
|
||||
|
||||
If you want to add new hardware, please contact us on [Slack](https://slack.vllm.ai/) or [Email](mailto:collaboration@vllm.ai).
|
||||
|
||||
@@ -31,6 +31,16 @@ vllm serve gs://core-llm/Llama-3-8b \
|
||||
--load-format runai_streamer
|
||||
```
|
||||
|
||||
To run model from Azure Blob Storage run:
|
||||
|
||||
```bash
|
||||
AZURE_STORAGE_ACCOUNT_NAME=<account> \
|
||||
vllm serve az://<container>/<model-path> \
|
||||
--load-format runai_streamer
|
||||
```
|
||||
|
||||
Authentication uses `DefaultAzureCredential`, which supports `az login`, managed identity, environment variables (`AZURE_CLIENT_ID`, `AZURE_TENANT_ID`, `AZURE_CLIENT_SECRET`), and other methods.
|
||||
|
||||
To run model from a S3 compatible object store run:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -418,6 +418,7 @@ th {
|
||||
| `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
|
||||
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
|
||||
| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
|
||||
| `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |
|
||||
| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `InternLM3ForCausalLM` | InternLM3 | `internlm/internlm3-8b-instruct`, etc. | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -21,7 +21,8 @@ vLLM provides multiple communication backends for EP. Use `--all2all-backend` to
|
||||
| `allgather_reducescatter` | Default backend | Standard all2all using allgather/reducescatter primitives | General purpose, works with any EP+DP configuration |
|
||||
| `deepep_high_throughput` | Multi-node prefill | Grouped GEMM with continuous layout, optimized for prefill | Prefill-dominated workloads, high-throughput scenarios |
|
||||
| `deepep_low_latency` | Multi-node decode | CUDA graph support, masked layout, optimized for decode | Decode-dominated workloads, low-latency scenarios |
|
||||
| `flashinfer_all2allv` | MNNVL systems | FlashInfer alltoallv kernels for multi-node NVLink | Systems with NVLink across nodes |
|
||||
| `flashinfer_nvlink_one_sided` | MNNVL systems | FlashInfer's one-sided A2A strategy for multi-node NVLink | High-throughput workloads |
|
||||
| `flashinfer_nvlink_two_sided` | MNNVL systems | FlashInfer's two-sided A2A strategy for multi-node NVLink | Systems with NVLink across nodes |
|
||||
| `naive` | Testing/debugging | Simple broadcast-based implementation | Debugging, not recommended for production |
|
||||
|
||||
## Single Node Deployment
|
||||
|
||||
@@ -72,6 +72,9 @@ In addition, we have the following custom APIs:
|
||||
- Only applicable to [classification models](../models/pooling_models.md).
|
||||
- [Score API](#score-api) (`/score`)
|
||||
- Applicable to [embedding models and cross-encoder models](../models/pooling_models.md).
|
||||
- [Cohere Embed API](#cohere-embed-api) (`/v2/embed`)
|
||||
- Compatible with [Cohere's Embed API](https://docs.cohere.com/reference/embed)
|
||||
- Works with any [embedding model](../models/pooling_models.md), including multimodal models.
|
||||
- [Re-rank API](#re-rank-api) (`/rerank`, `/v1/rerank`, `/v2/rerank`)
|
||||
- Implements [Jina AI's v1 re-rank API](https://jina.ai/reranker/)
|
||||
- Also compatible with [Cohere's v1 & v2 re-rank APIs](https://docs.cohere.com/v2/reference/rerank)
|
||||
@@ -429,6 +432,137 @@ these extra parameters are supported instead:
|
||||
--8<-- "vllm/entrypoints/pooling/base/protocol.py:embed-extra-params"
|
||||
```
|
||||
|
||||
### Cohere Embed API
|
||||
|
||||
Our API is also compatible with [Cohere's Embed v2 API](https://docs.cohere.com/reference/embed) which adds support for some modern embedding feature such as truncation, output dimensions, embedding types, and input types. This endpoint works with any embedding model (including multimodal models).
|
||||
|
||||
#### Cohere Embed API request parameters
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
| --------- | ---- | -------- | ----------- |
|
||||
| `model` | string | Yes | Model name |
|
||||
| `input_type` | string | No | Prompt prefix key (model-dependent, see below) |
|
||||
| `texts` | list[string] | No | Text inputs (use one of `texts`, `images`, or `inputs`) |
|
||||
| `images` | list[string] | No | Base64 data URI images |
|
||||
| `inputs` | list[object] | No | Mixed text and image content objects |
|
||||
| `embedding_types` | list[string] | No | Output types (default: `["float"]`) |
|
||||
| `output_dimension` | int | No | Truncate embeddings to this dimension (Matryoshka) |
|
||||
| `truncate` | string | No | `END`, `START`, or `NONE` (default: `END`) |
|
||||
|
||||
#### Text embedding
|
||||
|
||||
```bash
|
||||
curl -X POST "http://localhost:8000/v2/embed" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Snowflake/snowflake-arctic-embed-m-v1.5",
|
||||
"input_type": "query",
|
||||
"texts": ["Hello world", "How are you?"],
|
||||
"embedding_types": ["float"]
|
||||
}'
|
||||
```
|
||||
|
||||
??? console "Response"
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "embd-...",
|
||||
"embeddings": {
|
||||
"float": [
|
||||
[0.012, -0.034, ...],
|
||||
[0.056, 0.078, ...]
|
||||
]
|
||||
},
|
||||
"texts": ["Hello world", "How are you?"],
|
||||
"meta": {
|
||||
"api_version": {"version": "2"},
|
||||
"billed_units": {"input_tokens": 12}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Mixed text and image inputs
|
||||
|
||||
For multimodal models, you can embed images by passing base64 data URIs. The `inputs` field accepts a list of objects with mixed text and image content:
|
||||
|
||||
```bash
|
||||
curl -X POST "http://localhost:8000/v2/embed" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "google/siglip-so400m-patch14-384",
|
||||
"inputs": [
|
||||
{
|
||||
"content": [
|
||||
{"type": "text", "text": "A photo of a cat"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBOR..."}}
|
||||
]
|
||||
}
|
||||
],
|
||||
"embedding_types": ["float"]
|
||||
}'
|
||||
```
|
||||
|
||||
#### Embedding types
|
||||
|
||||
The `embedding_types` parameter controls the output format. Multiple types can be requested in a single call:
|
||||
|
||||
| Type | Description |
|
||||
| ---- | ----------- |
|
||||
| `float` | Raw float32 embeddings (default) |
|
||||
| `binary` | Bit-packed signed binary |
|
||||
| `ubinary` | Bit-packed unsigned binary |
|
||||
| `base64` | Little-endian float32 encoded as base64 |
|
||||
|
||||
```bash
|
||||
curl -X POST "http://localhost:8000/v2/embed" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Snowflake/snowflake-arctic-embed-m-v1.5",
|
||||
"input_type": "query",
|
||||
"texts": ["What is machine learning?"],
|
||||
"embedding_types": ["float", "binary"]
|
||||
}'
|
||||
```
|
||||
|
||||
??? console "Response"
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "embd-...",
|
||||
"embeddings": {
|
||||
"float": [[0.012, -0.034, ...]],
|
||||
"binary": [[42, -117, ...]]
|
||||
},
|
||||
"texts": ["What is machine learning?"],
|
||||
"meta": {
|
||||
"api_version": {"version": "2"},
|
||||
"billed_units": {"input_tokens": 8}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
#### Truncation
|
||||
|
||||
The `truncate` parameter controls how inputs exceeding the model's maximum sequence length are handled:
|
||||
|
||||
| Value | Behavior |
|
||||
| ----- | --------- |
|
||||
| `END` (default) | Keep the first tokens, drop the end |
|
||||
| `START` | Keep the last tokens, drop the beginning |
|
||||
| `NONE` | Return an error if the input is too long |
|
||||
|
||||
#### Input type and prompt prefixes
|
||||
|
||||
The `input_type` field selects a prompt prefix to prepend to each text input. The available values
|
||||
depend on the model:
|
||||
|
||||
- **Models with `task_instructions` in `config.json`**: The keys from the `task_instructions` dict are
|
||||
the valid `input_type` values and the corresponding value is prepended to each text.
|
||||
- **Models with `config_sentence_transformers.json` prompts**: The keys from the `prompts` dict are
|
||||
the valid `input_type` values. For example, `Snowflake/snowflake-arctic-embed-xs` defines `"query"`,
|
||||
so setting `input_type: "query"` prepends `"Represent this sentence for searching relevant passages: "`.
|
||||
- **Other models**: `input_type` is not accepted and will raise a validation error if passed.
|
||||
|
||||
### Transcriptions API
|
||||
|
||||
Our Transcriptions API is compatible with [OpenAI's Transcriptions API](https://platform.openai.com/docs/api-reference/audio/createTranscription);
|
||||
|
||||
@@ -14,6 +14,10 @@ import regex as re
|
||||
import zmq
|
||||
from quart import Quart, make_response, request
|
||||
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import (
|
||||
MoRIIOConstants,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.DEBUG)
|
||||
prefill_instances: list[dict] = []
|
||||
@@ -213,6 +217,8 @@ async def handle_request():
|
||||
|
||||
dip, dport = extract_ip_port_fast(decode_instance_endpoint["request_address"])
|
||||
|
||||
transfer_id = f"{MoRIIOConstants.TRANSFER_PREFIX}-{str(uuid.uuid4())}"
|
||||
|
||||
req_data_to_prefill = copy.deepcopy(req_data)
|
||||
req_data_to_prefill["kv_transfer_params"] = {}
|
||||
req_data["kv_transfer_params"] = {}
|
||||
@@ -222,6 +228,7 @@ async def handle_request():
|
||||
req_data_to_prefill["kv_transfer_params"]["remote_tp_size"] = (
|
||||
decode_instance_endpoint["tp_size"]
|
||||
)
|
||||
req_data_to_prefill["kv_transfer_params"]["transfer_id"] = transfer_id
|
||||
|
||||
send_prefill_task = asyncio.create_task(
|
||||
send_request_to_prefill(
|
||||
@@ -267,6 +274,7 @@ async def handle_request():
|
||||
|
||||
if selected_prefill_dp_rank is not None:
|
||||
req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank
|
||||
req_data["kv_transfer_params"]["transfer_id"] = transfer_id
|
||||
|
||||
decode_request_task = asyncio.create_task(
|
||||
start_decode_request(
|
||||
|
||||
@@ -24,7 +24,7 @@ outlines_core == 0.2.11
|
||||
# required for outlines backend disk cache
|
||||
diskcache == 5.6.3
|
||||
lark == 1.2.2
|
||||
xgrammar == 0.1.29; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
|
||||
xgrammar >= 0.1.32, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
|
||||
typing_extensions >= 4.10
|
||||
filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
|
||||
partial-json-parser # used for parsing partial JSON outputs
|
||||
|
||||
@@ -1,2 +1,2 @@
|
||||
# formatting
|
||||
pre-commit==4.0.1
|
||||
pre-commit>=4.5.1
|
||||
|
||||
@@ -42,7 +42,7 @@ tritonclient>=2.51.0
|
||||
|
||||
numba == 0.61.2 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs]==0.15.3
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.2.2
|
||||
instanttensor>=0.1.5
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
|
||||
@@ -45,6 +45,8 @@ pystemmer==3.0.0
|
||||
# via mteb
|
||||
|
||||
# Multi-modal processing
|
||||
av==16.1.0
|
||||
# required for audio_in_video tests
|
||||
blobfile==3.0.0
|
||||
# Multi-Modal Models Test
|
||||
decord==0.6.0
|
||||
|
||||
@@ -15,7 +15,7 @@ tensorizer==2.10.1
|
||||
packaging>=24.2
|
||||
setuptools>=77.0.3,<80.0.0
|
||||
setuptools-scm>=8
|
||||
runai-model-streamer[s3,gcs]==0.15.3
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
conch-triton-kernels==1.2.1
|
||||
timm>=1.0.17
|
||||
# amd-quark: required for Quark quantization on ROCm
|
||||
|
||||
@@ -10,6 +10,7 @@ pytest-cov
|
||||
|
||||
# testing utils
|
||||
albumentations # required for Nemotron Parse in test_common.py
|
||||
av # required for audio_in_video tests
|
||||
backoff # required for phi4mm test
|
||||
blobfile # required for kimi-vl test
|
||||
einops # required for MPT, qwen-vl
|
||||
@@ -55,7 +56,7 @@ grpcio-reflection==1.78.0
|
||||
arctic-inference == 0.1.1 # Required for suffix decoding test
|
||||
numba == 0.61.2 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs]==0.15.3
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
|
||||
instanttensor>=0.1.5
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
|
||||
+40
-5
@@ -62,6 +62,16 @@ attrs==24.2.0
|
||||
# referencing
|
||||
audioread==3.0.1
|
||||
# via librosa
|
||||
av==16.1.0
|
||||
# via -r requirements/test.in
|
||||
azure-core==1.38.2
|
||||
# via
|
||||
# azure-identity
|
||||
# azure-storage-blob
|
||||
azure-identity==1.25.2
|
||||
# via runai-model-streamer-azure
|
||||
azure-storage-blob==12.28.0
|
||||
# via runai-model-streamer-azure
|
||||
backoff==2.2.1
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
@@ -101,8 +111,10 @@ certifi==2024.8.30
|
||||
# rasterio
|
||||
# requests
|
||||
# sentry-sdk
|
||||
cffi==1.17.1
|
||||
# via soundfile
|
||||
cffi==2.0.0
|
||||
# via
|
||||
# cryptography
|
||||
# soundfile
|
||||
chardet==5.2.0
|
||||
# via mbstrdecoder
|
||||
charset-normalizer==3.4.0
|
||||
@@ -146,6 +158,12 @@ coverage==7.10.6
|
||||
# via pytest-cov
|
||||
cramjam==2.9.0
|
||||
# via fastparquet
|
||||
cryptography==46.0.5
|
||||
# via
|
||||
# azure-identity
|
||||
# azure-storage-blob
|
||||
# msal
|
||||
# pyjwt
|
||||
cuda-bindings==12.9.4
|
||||
# via torch
|
||||
cuda-pathfinder==1.3.3
|
||||
@@ -377,6 +395,8 @@ iniconfig==2.0.0
|
||||
# via pytest
|
||||
instanttensor==0.1.5
|
||||
# via -r requirements/test.in
|
||||
isodate==0.7.2
|
||||
# via azure-storage-blob
|
||||
isoduration==20.11.0
|
||||
# via jsonschema
|
||||
isort==5.13.2
|
||||
@@ -490,6 +510,12 @@ more-itertools==10.5.0
|
||||
# via lm-eval
|
||||
mpmath==1.3.0
|
||||
# via sympy
|
||||
msal==1.34.0
|
||||
# via
|
||||
# azure-identity
|
||||
# msal-extensions
|
||||
msal-extensions==1.3.1
|
||||
# via azure-identity
|
||||
msgpack==1.1.0
|
||||
# via
|
||||
# librosa
|
||||
@@ -826,6 +852,8 @@ pydantic-extra-types==2.10.5
|
||||
# via mistral-common
|
||||
pygments==2.18.0
|
||||
# via rich
|
||||
pyjwt==2.11.0
|
||||
# via msal
|
||||
pyogrio==0.11.0
|
||||
# via geopandas
|
||||
pyparsing==3.2.0
|
||||
@@ -943,6 +971,7 @@ regex==2024.9.11
|
||||
# transformers
|
||||
requests==2.32.3
|
||||
# via
|
||||
# azure-core
|
||||
# buildkite-test-collector
|
||||
# datasets
|
||||
# diffusers
|
||||
@@ -955,6 +984,7 @@ requests==2.32.3
|
||||
# lightly
|
||||
# lm-eval
|
||||
# mistral-common
|
||||
# msal
|
||||
# mteb
|
||||
# pooch
|
||||
# ray
|
||||
@@ -991,11 +1021,13 @@ rsa==4.9.1
|
||||
# via google-auth
|
||||
rtree==1.4.0
|
||||
# via torchgeo
|
||||
runai-model-streamer==0.15.3
|
||||
runai-model-streamer==0.15.7
|
||||
# via -r requirements/test.in
|
||||
runai-model-streamer-gcs==0.15.3
|
||||
runai-model-streamer-azure==0.15.7
|
||||
# via runai-model-streamer
|
||||
runai-model-streamer-s3==0.15.3
|
||||
runai-model-streamer-gcs==0.15.7
|
||||
# via runai-model-streamer
|
||||
runai-model-streamer-s3==0.15.7
|
||||
# via runai-model-streamer
|
||||
s3transfer==0.10.3
|
||||
# via boto3
|
||||
@@ -1264,6 +1296,9 @@ typing-extensions==4.15.0
|
||||
# aiosignal
|
||||
# albumentations
|
||||
# alembic
|
||||
# azure-core
|
||||
# azure-identity
|
||||
# azure-storage-blob
|
||||
# chz
|
||||
# fastapi
|
||||
# grpcio
|
||||
|
||||
@@ -657,13 +657,18 @@ class precompiled_wheel_utils:
|
||||
def get_base_commit_in_main_branch() -> str:
|
||||
try:
|
||||
# Get the latest commit hash of the upstream main branch.
|
||||
resp_json = subprocess.check_output(
|
||||
[
|
||||
"curl",
|
||||
"-s",
|
||||
"https://api.github.com/repos/vllm-project/vllm/commits/main",
|
||||
curl_cmd = [
|
||||
"curl",
|
||||
"-s",
|
||||
"https://api.github.com/repos/vllm-project/vllm/commits/main",
|
||||
]
|
||||
github_token = os.getenv("GH_TOKEN", os.getenv("GITHUB_TOKEN"))
|
||||
if github_token:
|
||||
curl_cmd += [
|
||||
"-H",
|
||||
f"Authorization: token {github_token}",
|
||||
]
|
||||
).decode("utf-8")
|
||||
resp_json = subprocess.check_output(curl_cmd).decode("utf-8")
|
||||
upstream_main_commit = json.loads(resp_json)["sha"]
|
||||
print(f"Upstream main branch latest commit: {upstream_main_commit}")
|
||||
|
||||
@@ -966,18 +971,19 @@ setup(
|
||||
ext_modules=ext_modules,
|
||||
install_requires=get_requirements(),
|
||||
extras_require={
|
||||
# AMD Zen CPU optimizations via zentorch
|
||||
"zen": ["zentorch"],
|
||||
"bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
|
||||
"tensorizer": ["tensorizer==2.10.1"],
|
||||
"fastsafetensors": ["fastsafetensors >= 0.2.2"],
|
||||
"instanttensor": ["instanttensor >= 0.1.5"],
|
||||
"runai": ["runai-model-streamer[s3,gcs] >= 0.15.3"],
|
||||
"runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
|
||||
"audio": [
|
||||
"librosa",
|
||||
"scipy",
|
||||
"soundfile",
|
||||
"mistral_common[audio]",
|
||||
"av",
|
||||
"torchcodec",
|
||||
], # Required for audio processing
|
||||
"video": [], # Kept for backwards compatibility
|
||||
"flashinfer": [], # Kept for backwards compatibility
|
||||
|
||||
@@ -179,7 +179,7 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
|
||||
def ops_in_model_before(self):
|
||||
return [
|
||||
torch.ops.vllm.all_reduce.default,
|
||||
torch.ops._C.scaled_fp4_quant.default,
|
||||
torch.ops._C.scaled_fp4_quant.out,
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -4,6 +4,9 @@
|
||||
|
||||
Tests the image source handling and tool_result content parsing in
|
||||
AnthropicServingMessages._convert_anthropic_to_openai_request().
|
||||
|
||||
Also covers extended-thinking edge cases such as ``redacted_thinking``
|
||||
blocks echoed back by Anthropic clients.
|
||||
"""
|
||||
|
||||
from vllm.entrypoints.anthropic.protocol import (
|
||||
@@ -373,3 +376,262 @@ class TestAttributionHeaderStripping:
|
||||
result = _convert(request)
|
||||
system_msg = result.messages[0]
|
||||
assert system_msg["content"] == "You are a helpful assistant."
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# Thinking block conversion (Anthropic → OpenAI)
|
||||
# ======================================================================
|
||||
|
||||
|
||||
class TestThinkingBlockConversion:
|
||||
"""Verify that thinking blocks in assistant messages are correctly
|
||||
moved to the ``reasoning`` field and stripped from ``content`` during
|
||||
the Anthropic→OpenAI conversion.
|
||||
|
||||
This is the Anthropic-endpoint path: the client echoes back the full
|
||||
assistant message (including thinking blocks emitted by vllm) in
|
||||
subsequent requests.
|
||||
"""
|
||||
|
||||
def test_thinking_plus_text_in_assistant_message(self):
|
||||
"""thinking + text → reasoning field + plain-string content."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Write me some code."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "I should write a simple example.",
|
||||
"signature": "sig_abc123",
|
||||
},
|
||||
{"type": "text", "text": "Sure! Here is the code."},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Can you fix the bug?"},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
# Find the assistant message in the converted output.
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
# Thinking content must be in reasoning, NOT in content.
|
||||
assert asst.get("reasoning") == "I should write a simple example."
|
||||
assert asst.get("content") == "Sure! Here is the code."
|
||||
|
||||
def test_thinking_only_in_assistant_message(self):
|
||||
"""Assistant message with only a thinking block (no visible text).
|
||||
|
||||
This can happen when the model emits reasoning but no final answer
|
||||
yet (e.g. a mid-turn reasoning step). Content should be None.
|
||||
"""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Hello"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "Just thinking...",
|
||||
"signature": "sig_xyz",
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Go on."},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
assert asst.get("reasoning") == "Just thinking..."
|
||||
# No visible text → content should be absent or None.
|
||||
assert asst.get("content") is None
|
||||
|
||||
def test_thinking_plus_tool_use_in_assistant_message(self):
|
||||
"""thinking + tool_use: reasoning field set, tool_calls populated."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "I need to call the calculator.",
|
||||
"signature": "sig_tool",
|
||||
},
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "call_001",
|
||||
"name": "calculator",
|
||||
"input": {"expression": "2+2"},
|
||||
},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "call_001",
|
||||
"content": "4",
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
assert asst.get("reasoning") == "I need to call the calculator."
|
||||
tool_calls = list(asst.get("tool_calls", []))
|
||||
assert len(tool_calls) == 1
|
||||
assert tool_calls[0]["function"]["name"] == "calculator"
|
||||
# No text content alongside reasoning + tool_use.
|
||||
assert asst.get("content") is None
|
||||
|
||||
def test_multiple_thinking_blocks_concatenated(self):
|
||||
"""Multiple thinking blocks should be joined in order."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Think hard."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "First thought. ",
|
||||
"signature": "s1",
|
||||
},
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "Second thought.",
|
||||
"signature": "s2",
|
||||
},
|
||||
{"type": "text", "text": "Done."},
|
||||
],
|
||||
},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
assert asst.get("reasoning") == "First thought. Second thought."
|
||||
assert asst.get("content") == "Done."
|
||||
|
||||
def test_no_thinking_blocks_unchanged(self):
|
||||
"""Messages without thinking blocks must not be modified."""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Hi"},
|
||||
{"role": "assistant", "content": "Hello!"},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
assert asst.get("content") == "Hello!"
|
||||
assert "reasoning" not in asst
|
||||
|
||||
def test_multi_turn_with_thinking_blocks(self):
|
||||
"""Full multi-turn conversation: previous assistant messages that
|
||||
include thinking blocks must all be converted without a 400 error.
|
||||
|
||||
This is the primary regression scenario from the bug report:
|
||||
upgrading vllm from v0.15.1 → v0.17.0 introduced thinking-block
|
||||
support in responses, but echoing those responses back in subsequent
|
||||
requests caused a Pydantic validation failure.
|
||||
"""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Turn 1 question"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "Reasoning for turn 1.",
|
||||
"signature": "s_t1",
|
||||
},
|
||||
{"type": "text", "text": "Answer for turn 1."},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Turn 2 question"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "Reasoning for turn 2.",
|
||||
"signature": "s_t2",
|
||||
},
|
||||
{"type": "text", "text": "Answer for turn 2."},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Turn 3 question"},
|
||||
]
|
||||
)
|
||||
# Must not raise a ValidationError / 400.
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 2
|
||||
|
||||
assert asst_msgs[0].get("reasoning") == "Reasoning for turn 1."
|
||||
assert asst_msgs[0].get("content") == "Answer for turn 1."
|
||||
assert asst_msgs[1].get("reasoning") == "Reasoning for turn 2."
|
||||
assert asst_msgs[1].get("content") == "Answer for turn 2."
|
||||
|
||||
def test_redacted_thinking_block_is_accepted(self):
|
||||
"""Anthropic clients may echo back redacted thinking blocks.
|
||||
|
||||
vLLM should accept these blocks (to avoid 400 validation errors)
|
||||
and ignore them when constructing the OpenAI-format prompt.
|
||||
"""
|
||||
request = _make_request(
|
||||
[
|
||||
{"role": "user", "content": "Hello"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "Thinking...",
|
||||
"signature": "sig_think",
|
||||
},
|
||||
{
|
||||
"type": "redacted_thinking",
|
||||
"data": "BASE64_OR_OTHER_OPAQUE_DATA",
|
||||
},
|
||||
{"type": "text", "text": "Hi!"},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Continue"},
|
||||
]
|
||||
)
|
||||
result = _convert(request)
|
||||
|
||||
asst_msgs = [m for m in result.messages if m.get("role") == "assistant"]
|
||||
assert len(asst_msgs) == 1
|
||||
asst = asst_msgs[0]
|
||||
|
||||
# Redacted thinking is ignored, normal thinking still becomes reasoning.
|
||||
assert asst.get("reasoning") == "Thinking..."
|
||||
assert asst.get("content") == "Hi!"
|
||||
|
||||
@@ -137,6 +137,59 @@ async def test_streaming_output_consistency(client: OpenAI, model_name: str):
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_streaming_logprobs(client: OpenAI, model_name: str):
|
||||
"""Test that streaming with logprobs returns valid logprob data on
|
||||
output_text.delta events and that top_logprobs has the requested count."""
|
||||
response = await client.responses.create(
|
||||
model=model_name,
|
||||
input="Say hello.",
|
||||
stream=True,
|
||||
top_logprobs=3,
|
||||
include=["message.output_text.logprobs"],
|
||||
)
|
||||
|
||||
events = []
|
||||
async for event in response:
|
||||
events.append(event)
|
||||
|
||||
assert len(events) > 0
|
||||
|
||||
# Collect all output_text.delta events that carry logprobs
|
||||
text_delta_events = [e for e in events if e.type == "response.output_text.delta"]
|
||||
assert len(text_delta_events) > 0, "Expected at least one text delta event"
|
||||
|
||||
for delta_event in text_delta_events:
|
||||
logprobs = delta_event.logprobs
|
||||
assert logprobs is not None, "logprobs should be present on text delta events"
|
||||
assert len(logprobs) > 0, "logprobs list should not be empty"
|
||||
for lp in logprobs:
|
||||
# Each logprob entry must have a token and a logprob value
|
||||
assert lp.token is not None
|
||||
assert isinstance(lp.logprob, float)
|
||||
assert lp.logprob <= 0.0, f"logprob should be <= 0, got {lp.logprob}"
|
||||
# top_logprobs should have up to 3 entries
|
||||
assert lp.top_logprobs is not None
|
||||
assert len(lp.top_logprobs) <= 3
|
||||
for tl in lp.top_logprobs:
|
||||
assert tl.token is not None
|
||||
assert isinstance(tl.logprob, float)
|
||||
|
||||
# Verify that top_logprobs are actually populated, not always empty
|
||||
all_top_logprobs = [
|
||||
tl for e in text_delta_events for lp in e.logprobs for tl in lp.top_logprobs
|
||||
]
|
||||
assert len(all_top_logprobs) > 0, (
|
||||
"Expected at least one top_logprobs entry across all delta events"
|
||||
)
|
||||
|
||||
# Verify the completed event still has valid output
|
||||
completed = events[-1]
|
||||
assert completed.type == "response.completed"
|
||||
assert completed.response.status == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_name", [MODEL_NAME])
|
||||
async def test_streaming_reasoning_tokens_e2e(client: OpenAI, model_name: str):
|
||||
|
||||
@@ -0,0 +1,175 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import base64
|
||||
import json
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from ...conftest import VideoTestAssets
|
||||
from ...utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen2.5-Omni-3B"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server():
|
||||
args = [
|
||||
"--max-model-len",
|
||||
"16384",
|
||||
"--enforce-eager",
|
||||
"--limit-mm-per-prompt",
|
||||
json.dumps({"audio": 3, "video": 3}),
|
||||
]
|
||||
|
||||
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.core_model
|
||||
@pytest.mark.asyncio
|
||||
async def test_online_audio_in_video(
|
||||
client: openai.AsyncOpenAI, video_assets: VideoTestAssets
|
||||
):
|
||||
"""Test video input with `audio_in_video=True`"""
|
||||
|
||||
# we don't use video_urls above because they missed audio stream.
|
||||
video_path = video_assets[0].video_path
|
||||
with open(video_path, "rb") as f:
|
||||
video_base64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What's in this video?"},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
# multi-turn to test mm processor cache as well
|
||||
for _ in range(2):
|
||||
chat_completion = await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
messages=messages,
|
||||
max_tokens=16,
|
||||
extra_body={
|
||||
"mm_processor_kwargs": {
|
||||
"use_audio_in_video": True,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert len(chat_completion.choices) == 1
|
||||
choice = chat_completion.choices[0]
|
||||
assert choice.finish_reason == "length"
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.asyncio
|
||||
async def test_online_audio_in_video_multi_videos(
|
||||
client: openai.AsyncOpenAI, video_assets: VideoTestAssets
|
||||
):
|
||||
"""Test multi-video input with `audio_in_video=True`"""
|
||||
|
||||
# we don't use video_urls above because they missed audio stream.
|
||||
video_path = video_assets[0].video_path
|
||||
with open(video_path, "rb") as f:
|
||||
video_base64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What's in these two videos?"},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
|
||||
},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
# multi-turn to test mm processor cache as well
|
||||
for _ in range(2):
|
||||
chat_completion = await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
messages=messages,
|
||||
max_tokens=16,
|
||||
extra_body={
|
||||
"mm_processor_kwargs": {
|
||||
"use_audio_in_video": True,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
assert len(chat_completion.choices) == 1
|
||||
choice = chat_completion.choices[0]
|
||||
assert choice.finish_reason == "length"
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.asyncio
|
||||
async def test_online_audio_in_video_interleaved(
|
||||
client: openai.AsyncOpenAI, video_assets: VideoTestAssets
|
||||
):
|
||||
"""Test interleaved video/audio input with `audio_in_video=True`"""
|
||||
|
||||
# we don't use video_urls above because they missed audio stream.
|
||||
video_path = video_assets[0].video_path
|
||||
with open(video_path, "rb") as f:
|
||||
video_base64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What's in these two videos?"},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": f"data:audio/mp4;base64,{video_base64}"},
|
||||
},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": f"data:video/mp4;base64,{video_base64}"},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
with pytest.raises(
|
||||
openai.BadRequestError,
|
||||
match="use_audio_in_video requires equal number of audio and video items",
|
||||
):
|
||||
await client.chat.completions.create(
|
||||
model=MODEL_NAME,
|
||||
messages=messages,
|
||||
max_tokens=16,
|
||||
extra_body={
|
||||
"mm_processor_kwargs": {
|
||||
"use_audio_in_video": True,
|
||||
}
|
||||
},
|
||||
)
|
||||
@@ -1,279 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Integration tests for GPT-OSS structural tags functionality (PR #25515)."""
|
||||
|
||||
import json
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.mcp.tool_server import ToolServer
|
||||
from vllm.reasoning.gptoss_reasoning_parser import (
|
||||
GptOssReasoningParser,
|
||||
)
|
||||
from vllm.sampling_params import StructuredOutputsParams
|
||||
|
||||
|
||||
class TestGptOssStructuralTagsIntegration:
|
||||
"""Integration tests for structural tags in GPT-OSS tool calls."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tokenizer(self):
|
||||
"""Create a mock tokenizer."""
|
||||
tokenizer = Mock()
|
||||
tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5])
|
||||
tokenizer.get_vocab = Mock(return_value={"<|end|>": 6})
|
||||
return tokenizer
|
||||
|
||||
@pytest.fixture
|
||||
def gptoss_parser(self, mock_tokenizer):
|
||||
"""Create a real GptOssReasoningParser instance."""
|
||||
return GptOssReasoningParser(mock_tokenizer)
|
||||
|
||||
@pytest.fixture
|
||||
def tool_server_with_python(self):
|
||||
"""Create a tool server with Python tool enabled."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=lambda tool: tool == "python")
|
||||
return tool_server
|
||||
|
||||
@pytest.fixture
|
||||
def tool_server_empty(self):
|
||||
"""Create a tool server with no tools."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(return_value=False)
|
||||
return tool_server
|
||||
|
||||
def test_end_to_end_no_tools(self, gptoss_parser):
|
||||
"""Test end-to-end flow when no tools are available."""
|
||||
# Test the parser directly
|
||||
result = gptoss_parser.prepare_structured_tag(None, None)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
# Verify basic structure
|
||||
assert parsed_result["type"] == "structural_tag"
|
||||
assert parsed_result["format"]["type"] == "triggered_tags"
|
||||
assert len(parsed_result["format"]["tags"]) == 1
|
||||
|
||||
# Verify only analysis channel is allowed
|
||||
analysis_tag = parsed_result["format"]["tags"][0]
|
||||
assert analysis_tag["begin"] == "<|channel|>analysis<|message|>"
|
||||
assert analysis_tag["content"]["type"] == "any_text"
|
||||
assert analysis_tag["end"] == "<|end|>"
|
||||
|
||||
# Verify triggers
|
||||
assert parsed_result["format"]["triggers"] == ["<|channel|>analysis"]
|
||||
assert parsed_result["format"]["stop_after_first"] is False
|
||||
|
||||
def test_end_to_end_with_python_tool(self, gptoss_parser, tool_server_with_python):
|
||||
"""Test end-to-end flow with Python tool enabled."""
|
||||
result = gptoss_parser.prepare_structured_tag(None, tool_server_with_python)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
# Should have analysis tag + 2 python tags
|
||||
assert len(parsed_result["format"]["tags"]) == 3
|
||||
|
||||
# Verify all expected tags are present
|
||||
tag_begins = [tag["begin"] for tag in parsed_result["format"]["tags"]]
|
||||
expected_begins = [
|
||||
"<|channel|>analysis<|message|>",
|
||||
"<|channel|>commentary to=python",
|
||||
"<|channel|>analysis to=python",
|
||||
]
|
||||
|
||||
for expected in expected_begins:
|
||||
assert expected in tag_begins
|
||||
|
||||
# Verify triggers include commentary
|
||||
assert "<|channel|>analysis" in parsed_result["format"]["triggers"]
|
||||
assert "<|channel|>commentary to=" in parsed_result["format"]["triggers"]
|
||||
|
||||
def test_structured_outputs_params_integration(
|
||||
self, gptoss_parser, tool_server_with_python
|
||||
):
|
||||
"""Test integration with StructuredOutputsParams."""
|
||||
# Generate structural tag
|
||||
structural_tag = gptoss_parser.prepare_structured_tag(
|
||||
None, tool_server_with_python
|
||||
)
|
||||
|
||||
# Create StructuredOutputsParams
|
||||
params = StructuredOutputsParams(structural_tag=structural_tag)
|
||||
|
||||
# Verify the tag is properly stored and accessible
|
||||
assert params.structural_tag == structural_tag
|
||||
|
||||
# Verify the tag is valid JSON
|
||||
parsed_tag = json.loads(params.structural_tag)
|
||||
assert parsed_tag["type"] == "structural_tag"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"browser, python, container, expected_tags",
|
||||
[
|
||||
# No tools
|
||||
(False, False, False, 1),
|
||||
# Single tool
|
||||
(True, False, False, 3),
|
||||
# Multiple tools
|
||||
(True, True, False, 5),
|
||||
# All tools
|
||||
(True, True, True, 7),
|
||||
],
|
||||
)
|
||||
def test_tool_server_interaction_flow(
|
||||
self, gptoss_parser, browser, python, container, expected_tags
|
||||
):
|
||||
"""Test the complete tool server interaction flow."""
|
||||
|
||||
# Create a mock ToolServer
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
|
||||
# Simulate tool availability based on parameters
|
||||
tool_server.has_tool = Mock(
|
||||
side_effect=lambda tool: {
|
||||
"browser": browser,
|
||||
"python": python,
|
||||
"container": container,
|
||||
}.get(tool, False)
|
||||
)
|
||||
|
||||
# Run the parser and verify results
|
||||
result = gptoss_parser.prepare_structured_tag(None, tool_server)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
# Validate number of tags
|
||||
assert len(parsed_result["format"]["tags"]) == expected_tags
|
||||
|
||||
# Verify tool-specific tags exist for enabled tools
|
||||
tag_begins = [tag["begin"] for tag in parsed_result["format"]["tags"]]
|
||||
for tool, enabled in {
|
||||
"browser": browser,
|
||||
"python": python,
|
||||
"container": container,
|
||||
}.items():
|
||||
if enabled:
|
||||
assert f"<|channel|>commentary to={tool}" in tag_begins
|
||||
assert f"<|channel|>analysis to={tool}" in tag_begins
|
||||
|
||||
def test_original_tag_preservation(self, gptoss_parser, tool_server_with_python):
|
||||
"""Test that original tags are preserved when provided."""
|
||||
original_tag = '{"type": "custom_tag", "data": "preserved"}'
|
||||
|
||||
result = gptoss_parser.prepare_structured_tag(
|
||||
original_tag, tool_server_with_python
|
||||
)
|
||||
|
||||
# Should return original tag unchanged
|
||||
assert result == original_tag
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tools",
|
||||
[
|
||||
[],
|
||||
["browser"],
|
||||
["python"],
|
||||
["container"],
|
||||
["browser", "python"],
|
||||
["browser", "container"],
|
||||
["python", "container"],
|
||||
["browser", "python", "container"],
|
||||
],
|
||||
)
|
||||
def test_json_validity_comprehensive(self, gptoss_parser, tools):
|
||||
"""Test JSON validity across all possible tool combinations."""
|
||||
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=lambda tool: tool in tools)
|
||||
|
||||
result = gptoss_parser.prepare_structured_tag(None, tool_server)
|
||||
|
||||
# Should be valid JSON
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
# Should have correct structure
|
||||
assert parsed_result["type"] == "structural_tag"
|
||||
assert "format" in parsed_result
|
||||
assert "tags" in parsed_result["format"]
|
||||
assert "triggers" in parsed_result["format"]
|
||||
|
||||
# Tag count should be: 1 (analysis) + 2 * len(tools)
|
||||
expected_tag_count = 1 + (2 * len(tools))
|
||||
assert len(parsed_result["format"]["tags"]) == expected_tag_count
|
||||
|
||||
def test_error_handling_invalid_tool_server(self, gptoss_parser):
|
||||
"""Test error handling with invalid tool server."""
|
||||
# Tool server that raises exceptions
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=Exception("Tool server error"))
|
||||
|
||||
# Should handle gracefully and still return a valid tag
|
||||
with pytest.raises(Exception, match="Tool server error"):
|
||||
gptoss_parser.prepare_structured_tag(None, tool_server)
|
||||
|
||||
def test_concurrent_requests_isolation(self, gptoss_parser):
|
||||
"""Test that concurrent requests don't interfere with each other."""
|
||||
# Simulate concurrent requests with different tool servers
|
||||
tool_server_1 = Mock(spec=ToolServer)
|
||||
tool_server_1.has_tool = Mock(side_effect=lambda tool: tool == "python")
|
||||
|
||||
tool_server_2 = Mock(spec=ToolServer)
|
||||
tool_server_2.has_tool = Mock(side_effect=lambda tool: tool == "browser")
|
||||
|
||||
# Generate tags concurrently
|
||||
result_1 = gptoss_parser.prepare_structured_tag(None, tool_server_1)
|
||||
result_2 = gptoss_parser.prepare_structured_tag(None, tool_server_2)
|
||||
|
||||
# Parse results
|
||||
parsed_1 = json.loads(result_1)
|
||||
parsed_2 = json.loads(result_2)
|
||||
|
||||
# Verify they have different tool configurations
|
||||
tags_1 = [tag["begin"] for tag in parsed_1["format"]["tags"]]
|
||||
tags_2 = [tag["begin"] for tag in parsed_2["format"]["tags"]]
|
||||
|
||||
# Result 1 should have python tags
|
||||
assert "<|channel|>commentary to=python" in tags_1
|
||||
assert "<|channel|>commentary to=browser" not in tags_1
|
||||
|
||||
# Result 2 should have browser tags
|
||||
assert "<|channel|>commentary to=browser" in tags_2
|
||||
assert "<|channel|>commentary to=python" not in tags_2
|
||||
|
||||
def test_tag_format_consistency(self, gptoss_parser):
|
||||
"""Test that all generated tags follow consistent format."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(
|
||||
side_effect=lambda tool: tool in ["python", "browser"]
|
||||
)
|
||||
|
||||
result = gptoss_parser.prepare_structured_tag(None, tool_server)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
# Verify all tags have required fields
|
||||
for tag in parsed_result["format"]["tags"]:
|
||||
assert "begin" in tag
|
||||
assert "content" in tag
|
||||
assert "end" in tag
|
||||
assert tag["content"]["type"] == "any_text"
|
||||
assert tag["end"] == "<|end|>"
|
||||
|
||||
# Verify begin format
|
||||
assert tag["begin"].startswith("<|channel|>")
|
||||
|
||||
def test_trigger_configuration(self, gptoss_parser):
|
||||
"""Test trigger configuration for different tool setups."""
|
||||
# Test with no tools
|
||||
result_no_tools = gptoss_parser.prepare_structured_tag(None, None)
|
||||
parsed_no_tools = json.loads(result_no_tools)
|
||||
assert parsed_no_tools["format"]["triggers"] == ["<|channel|>analysis"]
|
||||
|
||||
# Test with tools
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=lambda tool: tool == "python")
|
||||
|
||||
result_with_tools = gptoss_parser.prepare_structured_tag(None, tool_server)
|
||||
parsed_with_tools = json.loads(result_with_tools)
|
||||
|
||||
expected_triggers = ["<|channel|>analysis", "<|channel|>commentary to="]
|
||||
assert set(parsed_with_tools["format"]["triggers"]) == set(expected_triggers)
|
||||
@@ -0,0 +1,360 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import json
|
||||
import random
|
||||
from typing import Any
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.entrypoints.openai.engine.protocol import (
|
||||
DeltaMessage,
|
||||
)
|
||||
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
|
||||
MODEL = "ibm-granite/granite-4.0-h-tiny"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
model = MODEL
|
||||
args_for_model = [
|
||||
"--enforce-eager",
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"granite4",
|
||||
"--tokenizer",
|
||||
"ibm-granite/granite-4.0-h-tiny",
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
"--max-num-seqs",
|
||||
"2",
|
||||
]
|
||||
with RemoteOpenAIServer(model, args_for_model, max_wait_seconds=480) as server:
|
||||
yield server
|
||||
|
||||
|
||||
def create_complex_input(create_string_args: bool):
|
||||
coord_arg: dict | str = {
|
||||
"coordinates": [[23.54, 43.1], [-12.2, 54.3], [4, 5]],
|
||||
"coordinate_type": "latlong",
|
||||
}
|
||||
if create_string_args:
|
||||
# test granite behavior
|
||||
coord_arg = json.dumps(coord_arg)
|
||||
return [
|
||||
{"name": "find_bbox", "arguments": coord_arg},
|
||||
{
|
||||
"name": "get_stock_price",
|
||||
"arguments": {
|
||||
"symbol": "AAPL",
|
||||
"start_date": "2021-01-01",
|
||||
"end_date": "2021-12-31",
|
||||
},
|
||||
},
|
||||
{"name": "find_bbox", "arguments": coord_arg},
|
||||
]
|
||||
|
||||
|
||||
def random_chunks(s: str, min_len: int, max_len: int):
|
||||
chunks = []
|
||||
i = 0
|
||||
n = len(s)
|
||||
|
||||
while i < n:
|
||||
size = random.randint(min_len, max_len)
|
||||
chunks.append(s[i : i + size])
|
||||
i += size
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def tokenizer():
|
||||
return AutoTokenizer.from_pretrained(MODEL)
|
||||
|
||||
|
||||
# create a variety of input chunk sizes
|
||||
@pytest.mark.parametrize(
|
||||
"min_chunk, max_chunk",
|
||||
[
|
||||
(1, 1),
|
||||
(1, 2),
|
||||
(5, 7),
|
||||
(6, 20),
|
||||
],
|
||||
)
|
||||
def test_tool_call_parser_complex(min_chunk: int, max_chunk: int, tokenizer):
|
||||
input_dicts = create_complex_input(True)
|
||||
|
||||
formatted_tcs = [
|
||||
"<tool_call> " + json.dumps(call) + " </tool_call>" for call in input_dicts
|
||||
]
|
||||
|
||||
text_messages = [
|
||||
"Here goes the bbox call: \n",
|
||||
" Now the stock price call: \n ",
|
||||
" Now another bbox call: \n ",
|
||||
" See? I'm a helpful assistant.",
|
||||
]
|
||||
|
||||
test_input = (
|
||||
text_messages[0]
|
||||
+ formatted_tcs[0]
|
||||
+ text_messages[1]
|
||||
+ formatted_tcs[1]
|
||||
+ text_messages[2]
|
||||
+ formatted_tcs[2]
|
||||
+ text_messages[3]
|
||||
)
|
||||
|
||||
any_chat_request = ChatCompletionRequest(
|
||||
seed=42,
|
||||
model=MODEL,
|
||||
messages=[],
|
||||
)
|
||||
|
||||
parser = Granite4ToolParser(tokenizer=tokenizer)
|
||||
|
||||
delta_messages = list[DeltaMessage]()
|
||||
for text in random_chunks(test_input, min_chunk, max_chunk):
|
||||
delta = parser.extract_tool_calls_streaming(
|
||||
previous_text="",
|
||||
current_text="",
|
||||
delta_text=text,
|
||||
previous_token_ids=[],
|
||||
current_token_ids=[],
|
||||
delta_token_ids=[],
|
||||
request=any_chat_request,
|
||||
)
|
||||
if delta is not None:
|
||||
delta_messages.append(delta)
|
||||
|
||||
content = ""
|
||||
tool_calls = list[dict[str, Any]]()
|
||||
|
||||
current_name = "__start__"
|
||||
current_args = ""
|
||||
|
||||
for msg in delta_messages:
|
||||
if msg.content:
|
||||
content += msg.content
|
||||
for tool_call in msg.tool_calls:
|
||||
if delta_func := tool_call.function:
|
||||
if delta_func.name is not None:
|
||||
if current_name == "__start__":
|
||||
current_name = delta_func.name
|
||||
|
||||
if delta_func.name != current_name:
|
||||
tool_calls.append(
|
||||
{
|
||||
"name": current_name,
|
||||
"arguments": json.loads(current_args),
|
||||
}
|
||||
)
|
||||
current_name = delta_func.name
|
||||
current_args = ""
|
||||
|
||||
if delta_func.arguments:
|
||||
current_args += delta_func.arguments
|
||||
|
||||
if current_name != "__start__":
|
||||
tool_calls.append({"name": current_name, "arguments": json.loads(current_args)})
|
||||
|
||||
assert content == "".join(text_messages)
|
||||
assert tool_calls == create_complex_input(False)
|
||||
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_acme_region_name_for_transaction_id",
|
||||
"description": "Returns ACME transaction/transaction ID information"
|
||||
" including ACME regions\n\nArgs:\n start_time "
|
||||
"(str): Start date and time in datetime format "
|
||||
'"%Y-%m-%dT%H:%M:%S.%f"\n end_time (str): End '
|
||||
"date and time in datetime format "
|
||||
'"%Y-%m-%dT%H:%M:%S.%f"\n size (int, optional): '
|
||||
"Number of ACME Transaction IDs to return\n "
|
||||
"order (str, optional): Sort by most run "
|
||||
"transaction IDs. The value can be 'asc' for "
|
||||
"ascending or 'desc' for descending\n "
|
||||
"transaction_id (str, optional): ACME Transaction "
|
||||
"ID to filter on\n acme_region (str, optional): "
|
||||
"ACME Region to filter on\nReturns:\n - A "
|
||||
"dictionary containing a list of ACME transaction "
|
||||
"ids and the ACME regions they run in:\n {\n"
|
||||
' "Number of transaction IDs" : int,\n'
|
||||
' "Total transaction IDs available": int'
|
||||
',\n "ACME Transaction IDs": [\n '
|
||||
' {\n "Transaction ID": '
|
||||
'str,\n "Number of runs": int,\n'
|
||||
' "ACME Regions": [str],\n '
|
||||
" },\n ...\n ],"
|
||||
'\n "Start time" : datetime,\n '
|
||||
' "End time" : datetime,\n '
|
||||
' "Order" : str\n }\n '
|
||||
" - If no ACME region found for transaction id, "
|
||||
'returns:\n {"Success": "No ACME region '
|
||||
'found for transaction id."}\n - If an error '
|
||||
'occurs, returns:\n {"Error": "{exception'
|
||||
' message}"}',
|
||||
"parameters": {
|
||||
"properties": {
|
||||
"start_time": {},
|
||||
"end_time": {},
|
||||
"size": {"default": 500},
|
||||
"order": {"default": "desc"},
|
||||
"transaction_id": {"default": None},
|
||||
"acme_region": {"default": None},
|
||||
},
|
||||
"required": ["start_time", "end_time"],
|
||||
"type": "object",
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
tools2 = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"location": {
|
||||
"description": "The city and state, e.g. San Francisco, CA",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
"required": ["location"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_stock_price",
|
||||
"description": "Retrieves the current stock price for a given "
|
||||
"ticker symbol. The ticker symbol must be a valid "
|
||||
"symbol for a publicly traded company on a major US"
|
||||
" stock exchange like NYSE or NASDAQ. The tool will"
|
||||
" return the latest trade price in USD. It should "
|
||||
"be used when the user asks about the current or "
|
||||
"most recent price of a specific stock. It will not"
|
||||
" provide any other information about the stock or"
|
||||
" company.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"ticker": {
|
||||
"description": "The stock ticker symbol, e.g."
|
||||
" AAPL for Apple Inc.",
|
||||
"type": "string",
|
||||
}
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
messages = [
|
||||
{
|
||||
"content": "\n\nSystem: You are a helpful, precise, and methodical AI"
|
||||
" assistant that uses tool outputs provided inline.\nAlways"
|
||||
" assume the current datetime is 2026-01-29T13:59:09.238901"
|
||||
"+00:00.\n\nIf you receive a ToolMessage with `tool_call_id"
|
||||
'` equal to "get_time_range" (or "time_range_tool"), you '
|
||||
"MUST:\n 1. Parse that JSON and use the values `start` and"
|
||||
" `end` directly when calling other tools.\n 2. Do not "
|
||||
"re-call or re-compute the time range.\n 3. Pass resolved "
|
||||
"values (ISO strings) as arguments to any subsequent tool "
|
||||
"(do not pass function metadata or placeholders).\n 4. If "
|
||||
"a tool requires datetime objects rather than strings, "
|
||||
"convert the ISO strings into language-native datetime "
|
||||
"objects before invoking.\n\nAlways return fully resolved "
|
||||
"arguments in correct types (e.g., ISO datetime strings or"
|
||||
" datetime objects) and never include placeholders like "
|
||||
'"<start>".\n\n',
|
||||
"role": "system",
|
||||
},
|
||||
{
|
||||
"content": "What are the transaction IDs that ran in the"
|
||||
" ACME region A9345 over the last two months?",
|
||||
"role": "user",
|
||||
},
|
||||
{
|
||||
"content": '["2026-01-26T09: 51: 55.467722Z", "2026-01-27T09: 51: 55.467722Z"]',
|
||||
"role": "tool",
|
||||
"tool_call_id": "time_range_tool",
|
||||
},
|
||||
]
|
||||
messages2 = [{"role": "user", "content": "What's stock price for IBM?"}]
|
||||
|
||||
messages3 = [{"role": "user", "content": "What's the current weather in New York?"}]
|
||||
|
||||
|
||||
def get_args(client: openai.OpenAI, _tools, _messages, _stop):
|
||||
response = client.chat.completions.create(
|
||||
model=MODEL,
|
||||
messages=_messages,
|
||||
temperature=0,
|
||||
tools=_tools,
|
||||
max_tokens=200,
|
||||
stop=_stop,
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
return response.choices[0].message.tool_calls[0].function.arguments
|
||||
|
||||
|
||||
async def get_args_streaming(
|
||||
async_client: openai.AsyncOpenAI, _tools, _messages, _stop
|
||||
):
|
||||
stream = await async_client.chat.completions.create(
|
||||
model=MODEL,
|
||||
messages=_messages,
|
||||
temperature=0,
|
||||
tools=_tools,
|
||||
max_tokens=200,
|
||||
stop=_stop,
|
||||
tool_choice="auto",
|
||||
stream=True,
|
||||
)
|
||||
full_call = []
|
||||
async for chunk in stream:
|
||||
tc = chunk.choices[0].delta.tool_calls
|
||||
if tc and tc[0].function.arguments:
|
||||
full_call.append(tc[0].function.arguments)
|
||||
return "".join(full_call)
|
||||
|
||||
|
||||
async def run_scenario(server: RemoteOpenAIServer, _tools, _messages, _stop):
|
||||
non_streaming = get_args(server.get_client(), _tools, _messages, _stop)
|
||||
json.loads(non_streaming) # verify that it is json loadable
|
||||
streaming = await get_args_streaming(
|
||||
server.get_async_client(), _tools, _messages, _stop
|
||||
)
|
||||
json.loads(streaming)
|
||||
assert non_streaming == streaming, f"{non_streaming=}, {streaming=}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_sequence_interference(server: RemoteOpenAIServer):
|
||||
print("Testing scenario 1")
|
||||
await run_scenario(server, tools, messages, "veroniqueprattyushveroniqueprattyush")
|
||||
|
||||
print("Testing scenario 2")
|
||||
await run_scenario(
|
||||
server, tools2, messages2, "veroniqueprattyushveroniqueprattyush"
|
||||
)
|
||||
|
||||
print("Testing scenario 3")
|
||||
await run_scenario(server, tools2, messages3, "prattyush")
|
||||
@@ -3,29 +3,22 @@
|
||||
|
||||
import json
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from huggingface_hub import snapshot_download
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.tokenizers import TokenizerLike
|
||||
from vllm.tool_parsers.abstract_tool_parser import ToolParser
|
||||
from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser
|
||||
from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser
|
||||
|
||||
from ....utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
|
||||
LORA_MODEL = "minpeter/LoRA-Llama-3.2-1B-tool-vllm-ci"
|
||||
|
||||
SERVER_ARGS = [
|
||||
"--enforce-eager",
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"hermes",
|
||||
"--enable-lora",
|
||||
"--lora-modules",
|
||||
f"{LORA_MODEL}={LORA_MODEL}",
|
||||
"--tokenizer",
|
||||
f"{LORA_MODEL}",
|
||||
]
|
||||
|
||||
TOOLS = [
|
||||
{
|
||||
"type": "function",
|
||||
@@ -50,6 +43,75 @@ TOOLS = [
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
class ServerConfig(TypedDict, total=False):
|
||||
model: str
|
||||
arguments: list[str]
|
||||
model_arg: str
|
||||
tool_parser: ToolParser
|
||||
|
||||
|
||||
CONFIGS: dict[str, ServerConfig] = {
|
||||
"llama": {
|
||||
"model": "meta-llama/Llama-3.2-1B-Instruct",
|
||||
"arguments": [
|
||||
"--enforce-eager",
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"hermes",
|
||||
"--enable-lora",
|
||||
"--lora-modules",
|
||||
f"{LORA_MODEL}={LORA_MODEL}",
|
||||
"--tokenizer",
|
||||
f"{LORA_MODEL}",
|
||||
],
|
||||
"model_arg": LORA_MODEL,
|
||||
"tool_parser": Hermes2ProToolParser,
|
||||
},
|
||||
"granite4": {
|
||||
"model": "ibm-granite/granite-4.0-h-tiny",
|
||||
"arguments": [
|
||||
"--enforce-eager",
|
||||
"--enable-auto-tool-choice",
|
||||
"--tool-call-parser",
|
||||
"granite4",
|
||||
"--tokenizer",
|
||||
"ibm-granite/granite-4.0-h-tiny",
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
"--max-num-seqs",
|
||||
"2",
|
||||
],
|
||||
"model_arg": "ibm-granite/granite-4.0-h-tiny",
|
||||
"tool_parser": Granite4ToolParser,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# for each server config, download the model and return the config
|
||||
@pytest.fixture(scope="session", params=CONFIGS.keys())
|
||||
def server_config(request):
|
||||
config = CONFIGS[request.param]
|
||||
|
||||
# download model and tokenizer using transformers
|
||||
snapshot_download(config["model"])
|
||||
yield CONFIGS[request.param]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(request, server_config: ServerConfig):
|
||||
model = server_config["model"]
|
||||
args_for_model = server_config["arguments"]
|
||||
with RemoteOpenAIServer(model, args_for_model, max_wait_seconds=480) as server:
|
||||
yield server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def client(server: RemoteOpenAIServer):
|
||||
async with server.get_async_client() as async_client:
|
||||
yield async_client
|
||||
|
||||
|
||||
PRODUCT_TOOLS = [
|
||||
{
|
||||
"type": "function",
|
||||
@@ -87,186 +149,182 @@ PRODUCT_MESSAGES = [
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streaming_tool_call():
|
||||
async def test_non_streaming_tool_call(
|
||||
client: openai.AsyncOpenAI, server_config: ServerConfig
|
||||
):
|
||||
"""Test tool call in non-streaming mode."""
|
||||
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
|
||||
client = server.get_async_client()
|
||||
|
||||
response = await client.chat.completions.create(
|
||||
model=LORA_MODEL,
|
||||
messages=MESSAGES,
|
||||
tools=TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.0,
|
||||
)
|
||||
response = await client.chat.completions.create(
|
||||
model=server_config["model_arg"],
|
||||
messages=MESSAGES,
|
||||
tools=TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.0,
|
||||
)
|
||||
|
||||
assert response.choices
|
||||
choice = response.choices[0]
|
||||
message = choice.message
|
||||
assert response.choices
|
||||
choice = response.choices[0]
|
||||
message = choice.message
|
||||
|
||||
assert choice.finish_reason == "tool_calls"
|
||||
assert message.tool_calls is not None
|
||||
assert choice.finish_reason == "tool_calls"
|
||||
assert message.tool_calls is not None
|
||||
|
||||
tool_call = message.tool_calls[0]
|
||||
assert tool_call.type == "function"
|
||||
assert tool_call.function.name == "get_current_weather"
|
||||
tool_call = message.tool_calls[0]
|
||||
assert tool_call.type == "function"
|
||||
assert tool_call.function.name == "get_current_weather"
|
||||
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
assert "location" in arguments
|
||||
assert "Boston" in arguments["location"]
|
||||
print("\n[Non-Streaming Test Passed]")
|
||||
print(f"Tool Call: {tool_call.function.name}")
|
||||
print(f"Arguments: {arguments}")
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
assert "location" in arguments
|
||||
assert "Boston" in arguments["location"]
|
||||
print("\n[Non-Streaming Test Passed]")
|
||||
print(f"Tool Call: {tool_call.function.name}")
|
||||
print(f"Arguments: {arguments}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_tool_call():
|
||||
async def test_streaming_tool_call(
|
||||
client: openai.AsyncOpenAI, server_config: ServerConfig
|
||||
):
|
||||
"""Test tool call in streaming mode."""
|
||||
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
|
||||
client = server.get_async_client()
|
||||
|
||||
stream = await client.chat.completions.create(
|
||||
model=LORA_MODEL,
|
||||
messages=MESSAGES,
|
||||
tools=TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
)
|
||||
stream = await client.chat.completions.create(
|
||||
model=server_config["model_arg"],
|
||||
messages=MESSAGES,
|
||||
tools=TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
tool_call_chunks = {}
|
||||
async for chunk in stream:
|
||||
if not chunk.choices:
|
||||
continue
|
||||
tool_call_chunks = {}
|
||||
async for chunk in stream:
|
||||
if not chunk.choices:
|
||||
continue
|
||||
|
||||
delta = chunk.choices[0].delta
|
||||
if not delta or not delta.tool_calls:
|
||||
continue
|
||||
delta = chunk.choices[0].delta
|
||||
if not delta or not delta.tool_calls:
|
||||
continue
|
||||
|
||||
for tool_chunk in delta.tool_calls:
|
||||
index = tool_chunk.index
|
||||
if index not in tool_call_chunks:
|
||||
tool_call_chunks[index] = {"name": "", "arguments": ""}
|
||||
for tool_chunk in delta.tool_calls:
|
||||
index = tool_chunk.index
|
||||
if index not in tool_call_chunks:
|
||||
tool_call_chunks[index] = {"name": "", "arguments": ""}
|
||||
|
||||
if tool_chunk.function.name:
|
||||
tool_call_chunks[index]["name"] += tool_chunk.function.name
|
||||
if tool_chunk.function.arguments:
|
||||
tool_call_chunks[index]["arguments"] += (
|
||||
tool_chunk.function.arguments
|
||||
)
|
||||
if tool_chunk.function.name:
|
||||
tool_call_chunks[index]["name"] += tool_chunk.function.name
|
||||
if tool_chunk.function.arguments:
|
||||
tool_call_chunks[index]["arguments"] += tool_chunk.function.arguments
|
||||
|
||||
assert len(tool_call_chunks) == 1
|
||||
reconstructed_tool_call = tool_call_chunks[0]
|
||||
assert len(tool_call_chunks) == 1
|
||||
reconstructed_tool_call = tool_call_chunks[0]
|
||||
|
||||
assert reconstructed_tool_call["name"] == "get_current_weather"
|
||||
assert reconstructed_tool_call["name"] == "get_current_weather"
|
||||
|
||||
arguments = json.loads(reconstructed_tool_call["arguments"])
|
||||
assert "location" in arguments
|
||||
assert "Boston" in arguments["location"]
|
||||
print("\n[Streaming Test Passed]")
|
||||
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
|
||||
print(f"Reconstructed Arguments: {arguments}")
|
||||
arguments = json.loads(reconstructed_tool_call["arguments"])
|
||||
assert "location" in arguments
|
||||
assert "Boston" in arguments["location"]
|
||||
print("\n[Streaming Test Passed]")
|
||||
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
|
||||
print(f"Reconstructed Arguments: {arguments}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streaming_product_tool_call():
|
||||
async def test_non_streaming_product_tool_call(
|
||||
client: openai.AsyncOpenAI, server_config: ServerConfig
|
||||
):
|
||||
"""Test tool call integer and boolean parameters in non-streaming mode."""
|
||||
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
|
||||
client = server.get_async_client()
|
||||
|
||||
response = await client.chat.completions.create(
|
||||
model=LORA_MODEL,
|
||||
messages=PRODUCT_MESSAGES,
|
||||
tools=PRODUCT_TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.66,
|
||||
)
|
||||
response = await client.chat.completions.create(
|
||||
model=server_config["model_arg"],
|
||||
messages=PRODUCT_MESSAGES,
|
||||
tools=PRODUCT_TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.66,
|
||||
)
|
||||
|
||||
assert response.choices
|
||||
choice = response.choices[0]
|
||||
message = choice.message
|
||||
assert response.choices
|
||||
choice = response.choices[0]
|
||||
message = choice.message
|
||||
|
||||
assert choice.finish_reason == "tool_calls"
|
||||
assert message.tool_calls is not None
|
||||
assert choice.finish_reason == "tool_calls"
|
||||
assert message.tool_calls is not None
|
||||
|
||||
tool_call = message.tool_calls[0]
|
||||
assert tool_call.type == "function"
|
||||
assert tool_call.function.name == "get_product_info"
|
||||
tool_call = message.tool_calls[0]
|
||||
assert tool_call.type == "function"
|
||||
assert tool_call.function.name == "get_product_info"
|
||||
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
assert "product_id" in arguments
|
||||
assert "inserted" in arguments
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
assert "product_id" in arguments
|
||||
assert "inserted" in arguments
|
||||
|
||||
product_id = arguments.get("product_id")
|
||||
inserted = arguments.get("inserted")
|
||||
product_id = arguments.get("product_id")
|
||||
inserted = arguments.get("inserted")
|
||||
|
||||
assert isinstance(product_id, int)
|
||||
assert product_id == 7355608
|
||||
assert isinstance(inserted, bool)
|
||||
assert inserted is True
|
||||
assert isinstance(product_id, int)
|
||||
assert product_id == 7355608
|
||||
assert isinstance(inserted, bool)
|
||||
assert inserted is True
|
||||
|
||||
print("\n[Non-Streaming Product Test Passed]")
|
||||
print(f"Tool Call: {tool_call.function.name}")
|
||||
print(f"Arguments: {arguments}")
|
||||
print("\n[Non-Streaming Product Test Passed]")
|
||||
print(f"Tool Call: {tool_call.function.name}")
|
||||
print(f"Arguments: {arguments}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_product_tool_call():
|
||||
async def test_streaming_product_tool_call(
|
||||
client: openai.AsyncOpenAI, server_config: ServerConfig
|
||||
):
|
||||
"""Test tool call integer and boolean parameters in streaming mode."""
|
||||
with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server:
|
||||
client = server.get_async_client()
|
||||
|
||||
stream = await client.chat.completions.create(
|
||||
model=LORA_MODEL,
|
||||
messages=PRODUCT_MESSAGES,
|
||||
tools=PRODUCT_TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.66,
|
||||
stream=True,
|
||||
)
|
||||
stream = await client.chat.completions.create(
|
||||
model=server_config["model_arg"],
|
||||
messages=PRODUCT_MESSAGES,
|
||||
tools=PRODUCT_TOOLS,
|
||||
tool_choice="auto",
|
||||
temperature=0.66,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
tool_call_chunks = {}
|
||||
async for chunk in stream:
|
||||
if not chunk.choices:
|
||||
continue
|
||||
tool_call_chunks = {}
|
||||
async for chunk in stream:
|
||||
if not chunk.choices:
|
||||
continue
|
||||
|
||||
delta = chunk.choices[0].delta
|
||||
if not delta or not delta.tool_calls:
|
||||
continue
|
||||
delta = chunk.choices[0].delta
|
||||
if not delta or not delta.tool_calls:
|
||||
continue
|
||||
|
||||
for tool_chunk in delta.tool_calls:
|
||||
index = tool_chunk.index
|
||||
if index not in tool_call_chunks:
|
||||
tool_call_chunks[index] = {"name": "", "arguments": ""}
|
||||
for tool_chunk in delta.tool_calls:
|
||||
index = tool_chunk.index
|
||||
if index not in tool_call_chunks:
|
||||
tool_call_chunks[index] = {"name": "", "arguments": ""}
|
||||
|
||||
if tool_chunk.function.name:
|
||||
tool_call_chunks[index]["name"] += tool_chunk.function.name
|
||||
if tool_chunk.function.arguments:
|
||||
tool_call_chunks[index]["arguments"] += (
|
||||
tool_chunk.function.arguments
|
||||
)
|
||||
if tool_chunk.function.name:
|
||||
tool_call_chunks[index]["name"] += tool_chunk.function.name
|
||||
if tool_chunk.function.arguments:
|
||||
tool_call_chunks[index]["arguments"] += tool_chunk.function.arguments
|
||||
|
||||
assert len(tool_call_chunks) == 1
|
||||
reconstructed_tool_call = tool_call_chunks[0]
|
||||
assert len(tool_call_chunks) == 1
|
||||
reconstructed_tool_call = tool_call_chunks[0]
|
||||
|
||||
assert reconstructed_tool_call["name"] == "get_product_info"
|
||||
assert reconstructed_tool_call["name"] == "get_product_info"
|
||||
|
||||
arguments = json.loads(reconstructed_tool_call["arguments"])
|
||||
assert "product_id" in arguments
|
||||
assert "inserted" in arguments
|
||||
arguments = json.loads(reconstructed_tool_call["arguments"])
|
||||
assert "product_id" in arguments
|
||||
assert "inserted" in arguments
|
||||
|
||||
# Handle type coercion for streaming test as well
|
||||
product_id = arguments.get("product_id")
|
||||
inserted = arguments.get("inserted")
|
||||
# Handle type coercion for streaming test as well
|
||||
product_id = arguments.get("product_id")
|
||||
inserted = arguments.get("inserted")
|
||||
|
||||
assert isinstance(product_id, int)
|
||||
assert product_id == 7355608
|
||||
assert isinstance(inserted, bool)
|
||||
assert inserted is True
|
||||
assert isinstance(product_id, int)
|
||||
assert product_id == 7355608
|
||||
assert isinstance(inserted, bool)
|
||||
assert inserted is True
|
||||
|
||||
print("\n[Streaming Product Test Passed]")
|
||||
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
|
||||
print(f"Reconstructed Arguments: {arguments}")
|
||||
print("\n[Streaming Product Test Passed]")
|
||||
print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}")
|
||||
print(f"Reconstructed Arguments: {arguments}")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -276,9 +334,10 @@ def qwen_tokenizer() -> TokenizerLike:
|
||||
return get_tokenizer("Qwen/Qwen3-32B")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hermes_parser(qwen_tokenizer: TokenizerLike) -> Hermes2ProToolParser:
|
||||
return Hermes2ProToolParser(qwen_tokenizer)
|
||||
@pytest.fixture(params=CONFIGS.keys())
|
||||
def hermes_parser(request, qwen_tokenizer: TokenizerLike) -> ToolParser:
|
||||
config = CONFIGS[request.param]
|
||||
return config["tool_parser"](qwen_tokenizer)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
@@ -292,7 +351,7 @@ def any_chat_request() -> ChatCompletionRequest:
|
||||
|
||||
def test_hermes_parser_streaming_just_forward_text(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """This is some prior text that has nothing to do with tool calling."""
|
||||
@@ -324,7 +383,7 @@ def test_hermes_parser_streaming_just_forward_text(
|
||||
|
||||
def test_hermes_parser_streaming_failure_case_bug_19056(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """<tool_call>
|
||||
@@ -358,7 +417,7 @@ def test_hermes_parser_streaming_failure_case_bug_19056(
|
||||
|
||||
def test_hermes_parser_streaming(
|
||||
qwen_tokenizer: TokenizerLike,
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = '<tool_call>\
|
||||
@@ -387,16 +446,20 @@ def test_hermes_parser_streaming(
|
||||
delta_messages.append(delta)
|
||||
print(delta_messages)
|
||||
assert delta_messages[0].tool_calls[0].function.name == "get_current_temperature"
|
||||
tool_call_args = "".join(
|
||||
delta.tool_calls[0].function.arguments or "" for delta in delta_messages
|
||||
)
|
||||
assert tool_call_args == (
|
||||
'{"location":"San Francisco, California, United States", "unit": "celsius"}'
|
||||
# load to normalize whitespace
|
||||
tool_call_args = json.loads(
|
||||
"".join(
|
||||
delta.tool_calls[0].function.arguments or "" for delta in delta_messages
|
||||
)
|
||||
)
|
||||
assert tool_call_args == {
|
||||
"location": "San Francisco, California, United States",
|
||||
"unit": "celsius",
|
||||
}
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_no_tool_call(
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """This is not a tool call."""
|
||||
@@ -410,7 +473,7 @@ def test_hermes_parser_non_streaming_no_tool_call(
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_between_tags(
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
text = """<tool_call>
|
||||
@@ -428,9 +491,12 @@ def test_hermes_parser_non_streaming_tool_call_between_tags(
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_until_eos(
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
if isinstance(hermes_parser, Granite4ToolParser):
|
||||
pytest.skip(reason="The Granite4 tool parser enforces a complete response")
|
||||
|
||||
text = """<tool_call>
|
||||
{"name": "final_answer", "arguments": {"trigger": true}}"""
|
||||
tool_call = hermes_parser.extract_tool_calls(
|
||||
@@ -445,7 +511,7 @@ def test_hermes_parser_non_streaming_tool_call_until_eos(
|
||||
|
||||
|
||||
def test_hermes_parser_non_streaming_tool_call_invalid_json(
|
||||
hermes_parser: Hermes2ProToolParser,
|
||||
hermes_parser: ToolParser,
|
||||
any_chat_request: ChatCompletionRequest,
|
||||
) -> None:
|
||||
# Missing closing brace to trigger exception
|
||||
|
||||
@@ -0,0 +1,310 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the Cohere /v2/embed API with generic (non-Cohere) models.
|
||||
|
||||
Validates that the Cohere v2 embed endpoint works correctly with standard
|
||||
embedding models, covering text embedding, embedding type conversions,
|
||||
response structure, batching, normalisation, and semantic similarity.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
DTYPE = "bfloat16"
|
||||
|
||||
MODELS: list[tuple[str, list[str]]] = [
|
||||
("intfloat/multilingual-e5-small", []),
|
||||
(
|
||||
"Snowflake/snowflake-arctic-embed-m-v1.5",
|
||||
[
|
||||
"--trust_remote_code",
|
||||
"--hf_overrides",
|
||||
'{"matryoshka_dimensions":[256]}',
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=MODELS, ids=lambda m: m[0])
|
||||
def model_config(request):
|
||||
return request.param
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model_name(model_config):
|
||||
return model_config[0]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(model_config):
|
||||
name, extra_args = model_config
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
"--gpu-memory-utilization",
|
||||
"0.02",
|
||||
] + extra_args
|
||||
with RemoteOpenAIServer(name, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def _cohere_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
model_name: str,
|
||||
texts: list[str] | None = None,
|
||||
images: list[str] | None = None,
|
||||
input_type: str | None = None,
|
||||
embedding_types: list[str] | None = None,
|
||||
) -> dict:
|
||||
body: dict = {"model": model_name}
|
||||
if input_type is not None:
|
||||
body["input_type"] = input_type
|
||||
if texts is not None:
|
||||
body["texts"] = texts
|
||||
if images is not None:
|
||||
body["images"] = images
|
||||
if embedding_types is not None:
|
||||
body["embedding_types"] = embedding_types
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _openai_embed(
|
||||
server: RemoteOpenAIServer, model_name: str, texts: list[str]
|
||||
) -> dict:
|
||||
body = {"model": model_name, "input": texts, "encoding_format": "float"}
|
||||
resp = requests.post(server.url_for("/v1/embeddings"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _cosine_sim(a: list[float], b: list[float]) -> float:
|
||||
va, vb = np.array(a), np.array(b)
|
||||
return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb)))
|
||||
|
||||
|
||||
# -----------------------------------------------------------
|
||||
# Text embedding tests
|
||||
# -----------------------------------------------------------
|
||||
|
||||
|
||||
def test_basic_embed(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server, model_name, texts=["hello world"], embedding_types=["float"]
|
||||
)
|
||||
assert "embeddings" in r
|
||||
assert len(r["embeddings"]["float"]) == 1
|
||||
assert len(r["embeddings"]["float"][0]) > 0
|
||||
|
||||
|
||||
def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str):
|
||||
"""An input_type not defined in the model's prompt config should be
|
||||
rejected with a 400 error."""
|
||||
body = {
|
||||
"model": model_name,
|
||||
"input_type": "nonexistent_type",
|
||||
"texts": ["hello world"],
|
||||
"embedding_types": ["float"],
|
||||
}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
assert resp.status_code == 400
|
||||
assert "Unsupported input_type" in resp.json()["error"]["message"]
|
||||
|
||||
|
||||
def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str):
|
||||
"""Omitting input_type should always work (no prompt prefix applied)."""
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": ["hello world"],
|
||||
"embedding_types": ["float"],
|
||||
}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert len(data["embeddings"]["float"]) == 1
|
||||
|
||||
|
||||
def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str):
|
||||
"""v1 (OpenAI) and v2 (Cohere) endpoints should produce the same
|
||||
float embeddings for a generic model."""
|
||||
texts = ["hello world"]
|
||||
v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
|
||||
v1 = _openai_embed(server, model_name, texts)
|
||||
cos = _cosine_sim(v2["embeddings"]["float"][0], v1["data"][0]["embedding"])
|
||||
assert cos > 0.9999, f"v1/v2 parity failed, cosine={cos}"
|
||||
|
||||
|
||||
def test_embedding_types(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
model_name,
|
||||
texts=["test"],
|
||||
embedding_types=["float", "binary", "ubinary"],
|
||||
)
|
||||
dim = len(r["embeddings"]["float"][0])
|
||||
assert len(r["embeddings"]["binary"][0]) == dim // 8
|
||||
assert len(r["embeddings"]["ubinary"][0]) == dim // 8
|
||||
|
||||
|
||||
def test_response_structure(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(server, model_name, texts=["test"], embedding_types=["float"])
|
||||
assert "id" in r
|
||||
assert "embeddings" in r
|
||||
assert "texts" in r
|
||||
assert r["texts"] == ["test"]
|
||||
assert "meta" in r
|
||||
assert r["meta"]["api_version"]["version"] == "2"
|
||||
assert "billed_units" in r["meta"]
|
||||
assert r["meta"]["billed_units"]["input_tokens"] > 0
|
||||
assert r["meta"]["billed_units"]["image_tokens"] == 0
|
||||
|
||||
|
||||
def test_batch(server: RemoteOpenAIServer, model_name: str):
|
||||
texts = ["apple", "banana", "cherry"]
|
||||
r = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
|
||||
assert len(r["embeddings"]["float"]) == 3
|
||||
dim = len(r["embeddings"]["float"][0])
|
||||
for emb in r["embeddings"]["float"]:
|
||||
assert len(emb) == dim
|
||||
|
||||
|
||||
def test_l2_normalized(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server, model_name, texts=["hello world"], embedding_types=["float"]
|
||||
)
|
||||
emb = np.array(r["embeddings"]["float"][0])
|
||||
assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
|
||||
|
||||
|
||||
def test_semantic_similarity(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
model_name,
|
||||
texts=["machine learning", "deep learning", "chocolate cake recipe"],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
embs = r["embeddings"]["float"]
|
||||
cos_related = _cosine_sim(embs[0], embs[1])
|
||||
cos_unrelated = _cosine_sim(embs[0], embs[2])
|
||||
assert cos_related > cos_unrelated
|
||||
|
||||
|
||||
def test_missing_input_returns_error(server: RemoteOpenAIServer, model_name: str):
|
||||
body = {"model": model_name}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_base64_embedding_type(server: RemoteOpenAIServer, model_name: str):
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
model_name,
|
||||
texts=["test encoding"],
|
||||
embedding_types=["float", "base64"],
|
||||
)
|
||||
float_emb = r["embeddings"]["float"][0]
|
||||
b64_str = r["embeddings"]["base64"][0]
|
||||
decoded = struct.unpack(f"<{len(float_emb)}f", base64.b64decode(b64_str))
|
||||
np.testing.assert_allclose(float_emb, decoded, rtol=1e-5)
|
||||
|
||||
|
||||
# -----------------------------------------------------------
|
||||
# Truncation tests
|
||||
# -----------------------------------------------------------
|
||||
|
||||
|
||||
def _cohere_embed_raw(
|
||||
server: RemoteOpenAIServer,
|
||||
body: dict,
|
||||
) -> requests.Response:
|
||||
return requests.post(server.url_for("/v2/embed"), json=body)
|
||||
|
||||
|
||||
def test_truncate_end_succeeds(server: RemoteOpenAIServer, model_name: str):
|
||||
"""truncate=END should silently truncate long input."""
|
||||
long_text = " ".join(["word"] * 2000)
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "END",
|
||||
}
|
||||
resp = _cohere_embed_raw(server, body)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert len(data["embeddings"]["float"]) == 1
|
||||
|
||||
|
||||
def test_truncate_start_succeeds(server: RemoteOpenAIServer, model_name: str):
|
||||
"""truncate=START should silently truncate long input from the start."""
|
||||
long_text = " ".join(["word"] * 2000)
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "START",
|
||||
}
|
||||
resp = _cohere_embed_raw(server, body)
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert len(data["embeddings"]["float"]) == 1
|
||||
|
||||
|
||||
def test_truncate_none_rejects_long_input(server: RemoteOpenAIServer, model_name: str):
|
||||
"""truncate=NONE should error when input exceeds model context."""
|
||||
long_text = " ".join(["word"] * 2000)
|
||||
body = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "NONE",
|
||||
}
|
||||
resp = _cohere_embed_raw(server, body)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
def test_truncate_start_vs_end_differ(server: RemoteOpenAIServer, model_name: str):
|
||||
"""START and END truncation should produce different embeddings
|
||||
when the input is long enough to actually be truncated.
|
||||
|
||||
We construct input with distinct tokens at the start vs end
|
||||
so that keeping different halves produces different embeddings.
|
||||
"""
|
||||
start_words = " ".join([f"alpha{i}" for i in range(300)])
|
||||
end_words = " ".join([f"omega{i}" for i in range(300)])
|
||||
long_text = start_words + " " + end_words
|
||||
|
||||
body_end = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "END",
|
||||
}
|
||||
body_start = {
|
||||
"model": model_name,
|
||||
"texts": [long_text],
|
||||
"embedding_types": ["float"],
|
||||
"truncate": "START",
|
||||
}
|
||||
r_end = _cohere_embed_raw(server, body_end).json()
|
||||
r_start = _cohere_embed_raw(server, body_start).json()
|
||||
|
||||
emb_end = r_end["embeddings"]["float"][0]
|
||||
emb_start = r_start["embeddings"]["float"][0]
|
||||
cos = _cosine_sim(emb_end, emb_start)
|
||||
assert cos < 0.99, (
|
||||
f"START and END truncation should produce different embeddings "
|
||||
f"for long input, but cosine similarity was {cos}"
|
||||
)
|
||||
@@ -0,0 +1,135 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the Cohere /v2/embed API with a multimodal model (SigLIP).
|
||||
|
||||
Validates image embedding, batching, normalisation, and embedding type
|
||||
conversions through the /v2/embed endpoint.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import struct
|
||||
import zlib
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "google/siglip-so400m-patch14-384"
|
||||
DTYPE = "bfloat16"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"64",
|
||||
"--gpu-memory-utilization",
|
||||
"0.3",
|
||||
]
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def _make_tiny_png(r: int, g: int, b: int, w: int = 2, h: int = 2) -> str:
|
||||
raw = b""
|
||||
for _ in range(h):
|
||||
raw += b"\x00" + bytes([r, g, b]) * w
|
||||
compressed = zlib.compress(raw)
|
||||
|
||||
def chunk(ctype: bytes, cdata: bytes) -> bytes:
|
||||
c = ctype + cdata
|
||||
return (
|
||||
struct.pack(">I", len(cdata))
|
||||
+ c
|
||||
+ struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)
|
||||
)
|
||||
|
||||
ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
|
||||
png = (
|
||||
b"\x89PNG\r\n\x1a\n"
|
||||
+ chunk(b"IHDR", ihdr)
|
||||
+ chunk(b"IDAT", compressed)
|
||||
+ chunk(b"IEND", b"")
|
||||
)
|
||||
return "data:image/png;base64," + base64.b64encode(png).decode()
|
||||
|
||||
|
||||
def _cohere_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
texts: list[str] | None = None,
|
||||
images: list[str] | None = None,
|
||||
embedding_types: list[str] | None = None,
|
||||
) -> dict:
|
||||
body: dict = {"model": MODEL_NAME}
|
||||
if texts is not None:
|
||||
body["texts"] = texts
|
||||
if images is not None:
|
||||
body["images"] = images
|
||||
if embedding_types is not None:
|
||||
body["embedding_types"] = embedding_types
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
def test_image_embed(server: RemoteOpenAIServer):
|
||||
img_uri = _make_tiny_png(255, 0, 0)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[img_uri],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
assert "embeddings" in r
|
||||
assert len(r["embeddings"]["float"]) == 1
|
||||
assert len(r["embeddings"]["float"][0]) > 0
|
||||
assert r["meta"]["billed_units"]["image_tokens"] > 0
|
||||
assert r["meta"]["billed_units"]["input_tokens"] == 0
|
||||
|
||||
|
||||
def test_image_batch(server: RemoteOpenAIServer):
|
||||
red = _make_tiny_png(255, 0, 0)
|
||||
blue = _make_tiny_png(0, 0, 255)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[red, blue],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
assert len(r["embeddings"]["float"]) == 2
|
||||
|
||||
|
||||
def test_image_l2_normalized(server: RemoteOpenAIServer):
|
||||
img_uri = _make_tiny_png(0, 255, 0)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[img_uri],
|
||||
embedding_types=["float"],
|
||||
)
|
||||
emb = np.array(r["embeddings"]["float"][0])
|
||||
assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
|
||||
|
||||
|
||||
def test_image_embedding_types(server: RemoteOpenAIServer):
|
||||
img_uri = _make_tiny_png(128, 128, 128)
|
||||
r = _cohere_embed(
|
||||
server,
|
||||
images=[img_uri],
|
||||
embedding_types=["float", "binary", "ubinary"],
|
||||
)
|
||||
dim = len(r["embeddings"]["float"][0])
|
||||
assert len(r["embeddings"]["binary"][0]) == dim // 8
|
||||
assert len(r["embeddings"]["ubinary"][0]) == dim // 8
|
||||
|
||||
|
||||
def test_text_embed_on_multimodal(server: RemoteOpenAIServer):
|
||||
"""SigLIP also supports text-only embedding via /v2/embed."""
|
||||
r = _cohere_embed(server, texts=["hello world"], embedding_types=["float"])
|
||||
assert "embeddings" in r
|
||||
assert len(r["embeddings"]["float"]) == 1
|
||||
assert len(r["embeddings"]["float"][0]) > 0
|
||||
@@ -0,0 +1,102 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Parity test between Cohere /v2/embed and OpenAI /v1/embeddings.
|
||||
|
||||
Verifies that both endpoints produce identical float embeddings when
|
||||
no prompt prefix is applied (input_type omitted for Cohere /v2/embed).
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
MODEL_NAME = "BAAI/bge-base-en-v1.5"
|
||||
DTYPE = "bfloat16"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server():
|
||||
args = [
|
||||
"--runner",
|
||||
"pooling",
|
||||
"--dtype",
|
||||
DTYPE,
|
||||
"--enforce-eager",
|
||||
"--max-model-len",
|
||||
"512",
|
||||
"--gpu-memory-utilization",
|
||||
"0.02",
|
||||
]
|
||||
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
def _cohere_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
texts: list[str],
|
||||
) -> list[list[float]]:
|
||||
body = {
|
||||
"model": MODEL_NAME,
|
||||
"texts": texts,
|
||||
"embedding_types": ["float"],
|
||||
}
|
||||
resp = requests.post(server.url_for("/v2/embed"), json=body)
|
||||
resp.raise_for_status()
|
||||
return resp.json()["embeddings"]["float"]
|
||||
|
||||
|
||||
def _openai_embed(
|
||||
server: RemoteOpenAIServer,
|
||||
texts: list[str],
|
||||
) -> list[list[float]]:
|
||||
body = {"model": MODEL_NAME, "input": texts, "encoding_format": "float"}
|
||||
resp = requests.post(server.url_for("/v1/embeddings"), json=body)
|
||||
resp.raise_for_status()
|
||||
return [item["embedding"] for item in resp.json()["data"]]
|
||||
|
||||
|
||||
def test_single_text_parity(server: RemoteOpenAIServer):
|
||||
"""A single text should produce identical embeddings via both APIs."""
|
||||
texts = ["the quick brown fox jumps over the lazy dog"]
|
||||
v2 = _cohere_embed(server, texts)
|
||||
v1 = _openai_embed(server, texts)
|
||||
np.testing.assert_allclose(v2[0], v1[0], rtol=1e-5)
|
||||
|
||||
|
||||
def test_batch_parity(server: RemoteOpenAIServer):
|
||||
"""A batch of texts should produce identical embeddings via both APIs,
|
||||
in the same order."""
|
||||
texts = [
|
||||
"machine learning",
|
||||
"deep learning",
|
||||
"natural language processing",
|
||||
]
|
||||
v2 = _cohere_embed(server, texts)
|
||||
v1 = _openai_embed(server, texts)
|
||||
assert len(v2) == len(v1) == 3
|
||||
for i in range(3):
|
||||
np.testing.assert_allclose(v2[i], v1[i], rtol=1e-5, err_msg=f"index {i}")
|
||||
|
||||
|
||||
def test_token_count_parity(server: RemoteOpenAIServer):
|
||||
"""Both APIs should report the same prompt token count."""
|
||||
texts = ["hello world"]
|
||||
v2_resp = requests.post(
|
||||
server.url_for("/v2/embed"),
|
||||
json={
|
||||
"model": MODEL_NAME,
|
||||
"texts": texts,
|
||||
"embedding_types": ["float"],
|
||||
},
|
||||
)
|
||||
v1_resp = requests.post(
|
||||
server.url_for("/v1/embeddings"),
|
||||
json={"model": MODEL_NAME, "input": texts, "encoding_format": "float"},
|
||||
)
|
||||
v2_resp.raise_for_status()
|
||||
v1_resp.raise_for_status()
|
||||
v2_tokens = v2_resp.json()["meta"]["billed_units"]["input_tokens"]
|
||||
v1_tokens = v1_resp.json()["usage"]["prompt_tokens"]
|
||||
assert v2_tokens == v1_tokens
|
||||
@@ -0,0 +1,208 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for EmbedIOProcessor."""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
|
||||
from vllm.entrypoints.pooling.embed.protocol import (
|
||||
CohereEmbedRequest,
|
||||
)
|
||||
|
||||
|
||||
class TestResolveTruncation:
|
||||
"""Unit tests for EmbedIOProcessor._resolve_cohere_truncation."""
|
||||
|
||||
@staticmethod
|
||||
def _make_request(**kwargs) -> CohereEmbedRequest:
|
||||
defaults = {
|
||||
"model": "test",
|
||||
"input_type": "search_document",
|
||||
"texts": ["hello"],
|
||||
}
|
||||
return CohereEmbedRequest(**(defaults | kwargs))
|
||||
|
||||
def test_truncate_end_default(self):
|
||||
req = self._make_request()
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == -1
|
||||
assert side is None
|
||||
|
||||
def test_truncate_end_explicit(self):
|
||||
req = self._make_request(truncate="END")
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == -1
|
||||
assert side is None
|
||||
|
||||
def test_truncate_end_with_max_tokens(self):
|
||||
req = self._make_request(truncate="END", max_tokens=128)
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == 128
|
||||
assert side is None
|
||||
|
||||
def test_truncate_none(self):
|
||||
req = self._make_request(truncate="NONE")
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens is None
|
||||
assert side is None
|
||||
|
||||
def test_truncate_none_with_max_tokens(self):
|
||||
"""truncate=NONE should NOT set truncate_prompt_tokens; the
|
||||
max_tokens limit is enforced separately via _check_max_tokens."""
|
||||
req = self._make_request(truncate="NONE", max_tokens=10)
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens is None
|
||||
assert side is None
|
||||
|
||||
def test_truncate_start(self):
|
||||
req = self._make_request(truncate="START")
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == -1
|
||||
assert side == "left"
|
||||
|
||||
def test_truncate_start_with_max_tokens(self):
|
||||
req = self._make_request(truncate="START", max_tokens=64)
|
||||
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
|
||||
assert tokens == 64
|
||||
assert side == "left"
|
||||
|
||||
|
||||
class TestApplyStPrompt:
|
||||
"""Unit tests for EmbedIOProcessor._apply_task_instruction."""
|
||||
|
||||
@staticmethod
|
||||
def _make_handler(task_instructions: dict[str, str] | None):
|
||||
handler = object.__new__(EmbedIOProcessor)
|
||||
handler.task_instructions = task_instructions
|
||||
return handler
|
||||
|
||||
def test_no_prompts_configured(self):
|
||||
handler = self._make_handler(None)
|
||||
texts = ["hello", "world"]
|
||||
assert handler._apply_task_instruction(texts, "query") is texts
|
||||
|
||||
def test_matching_input_type(self):
|
||||
handler = self._make_handler({"query": "search_query: "})
|
||||
result = handler._apply_task_instruction(["hello"], "query")
|
||||
assert result == ["search_query: hello"]
|
||||
|
||||
def test_non_matching_input_type(self):
|
||||
handler = self._make_handler({"query": "search_query: "})
|
||||
texts = ["hello"]
|
||||
assert handler._apply_task_instruction(texts, "document") is texts
|
||||
|
||||
def test_multiple_texts(self):
|
||||
handler = self._make_handler(
|
||||
{"query": "Represent this sentence for searching: "}
|
||||
)
|
||||
result = handler._apply_task_instruction(["a", "b", "c"], "query")
|
||||
assert result == [
|
||||
"Represent this sentence for searching: a",
|
||||
"Represent this sentence for searching: b",
|
||||
"Represent this sentence for searching: c",
|
||||
]
|
||||
|
||||
def test_empty_prefix_returns_unchanged(self):
|
||||
handler = self._make_handler({"passage": ""})
|
||||
texts = ["hello"]
|
||||
assert handler._apply_task_instruction(texts, "passage") is texts
|
||||
|
||||
|
||||
class TestLoadTaskInstructions:
|
||||
"""Unit tests for EmbedIOProcessor._load_task_instructions."""
|
||||
|
||||
def test_no_attribute(self):
|
||||
class FakeConfig:
|
||||
pass
|
||||
|
||||
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
|
||||
|
||||
def test_with_task_instructions(self):
|
||||
class FakeConfig:
|
||||
task_instructions = {
|
||||
"retrieval.query": "Represent the query: ",
|
||||
"retrieval.passage": "",
|
||||
}
|
||||
|
||||
result = EmbedIOProcessor._load_task_instructions(FakeConfig())
|
||||
assert result == {
|
||||
"retrieval.query": "Represent the query: ",
|
||||
"retrieval.passage": "",
|
||||
}
|
||||
|
||||
def test_empty_dict(self):
|
||||
class FakeConfig:
|
||||
task_instructions = {}
|
||||
|
||||
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
|
||||
|
||||
def test_non_dict(self):
|
||||
class FakeConfig:
|
||||
task_instructions = "not a dict"
|
||||
|
||||
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
|
||||
|
||||
|
||||
class TestCheckMaxTokens:
|
||||
"""Unit tests for EmbedIOProcessor._check_cohere_max_tokens."""
|
||||
|
||||
@staticmethod
|
||||
def _fake_output(n_tokens: int):
|
||||
class _Out:
|
||||
def __init__(self, n: int):
|
||||
self.prompt_token_ids = list(range(n))
|
||||
|
||||
return _Out(n_tokens)
|
||||
|
||||
def test_none_check_is_noop(self):
|
||||
outs = [self._fake_output(100)]
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, None)
|
||||
|
||||
def test_within_limit(self):
|
||||
outs = [self._fake_output(5), self._fake_output(3)]
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
|
||||
|
||||
def test_exceeds_limit(self):
|
||||
outs = [self._fake_output(3), self._fake_output(10)]
|
||||
with pytest.raises(ValueError, match="exceeds max_tokens=5"):
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
|
||||
|
||||
def test_exact_limit(self):
|
||||
outs = [self._fake_output(5)]
|
||||
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
|
||||
|
||||
|
||||
class TestValidateInputType:
|
||||
"""Unit tests for EmbedIOProcessor._validate_input_type."""
|
||||
|
||||
@staticmethod
|
||||
def _make_handler(task_instructions: dict[str, str] | None):
|
||||
handler = object.__new__(EmbedIOProcessor)
|
||||
handler.task_instructions = task_instructions
|
||||
return handler
|
||||
|
||||
def test_none_input_type_always_accepted(self):
|
||||
handler = self._make_handler(None)
|
||||
handler._validate_input_type(None)
|
||||
handler_with = self._make_handler({"query": "q: "})
|
||||
handler_with._validate_input_type(None)
|
||||
|
||||
def test_no_prompts_rejects(self):
|
||||
handler = self._make_handler(None)
|
||||
with pytest.raises(ValueError, match="does not define any input_type"):
|
||||
handler._validate_input_type("anything")
|
||||
|
||||
def test_known_type_accepted(self):
|
||||
handler = self._make_handler({"query": "q: ", "document": "d: "})
|
||||
handler._validate_input_type("query")
|
||||
handler._validate_input_type("document")
|
||||
|
||||
def test_unknown_type_rejected(self):
|
||||
handler = self._make_handler({"query": "q: ", "document": "d: "})
|
||||
with pytest.raises(ValueError, match="Unsupported input_type 'other'"):
|
||||
handler._validate_input_type("other")
|
||||
|
||||
def test_error_lists_supported(self):
|
||||
handler = self._make_handler({"a": "", "b": ""})
|
||||
with pytest.raises(ValueError, match="Supported values: a, b"):
|
||||
handler._validate_input_type("z")
|
||||
@@ -0,0 +1,129 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for Cohere embed protocol: build_typed_embeddings and its
|
||||
underlying packing helpers, plus Cohere-specific serving helpers."""
|
||||
|
||||
import base64
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.pooling.embed.protocol import (
|
||||
build_typed_embeddings,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_embeddings() -> list[list[float]]:
|
||||
return [
|
||||
[0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, -0.8],
|
||||
[-0.05, 0.15, -0.25, 0.35, -0.45, 0.55, -0.65, 0.75],
|
||||
]
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsFloat:
|
||||
def test_float_passthrough(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["float"])
|
||||
assert result.float == sample_embeddings
|
||||
assert result.binary is None
|
||||
|
||||
def test_empty_input(self):
|
||||
result = build_typed_embeddings([], ["float"])
|
||||
assert result.float == []
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsBinary:
|
||||
def test_binary_packing(self):
|
||||
# 8 values: positive->1, negative->0 => bits: 10101010 = 0xAA = 170
|
||||
# signed: 170 - 128 = 42
|
||||
embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
assert result.binary[0] == [42]
|
||||
|
||||
def test_ubinary_packing(self):
|
||||
embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
|
||||
result = build_typed_embeddings(embs, ["ubinary"])
|
||||
assert result.ubinary is not None
|
||||
assert result.ubinary[0] == [170] # 0b10101010
|
||||
|
||||
def test_binary_all_positive(self):
|
||||
embs = [[0.1] * 8]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
# all bits = 1 => 0xFF = 255, signed: 255 - 128 = 127
|
||||
assert result.binary[0] == [127]
|
||||
|
||||
def test_binary_all_negative(self):
|
||||
embs = [[-0.1] * 8]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
# all bits = 0, signed: 0 - 128 = -128
|
||||
assert result.binary[0] == [-128]
|
||||
|
||||
def test_binary_dimension_is_eighth(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["binary"])
|
||||
assert result.binary is not None
|
||||
for orig, packed in zip(sample_embeddings, result.binary):
|
||||
assert len(packed) == len(orig) // 8
|
||||
|
||||
def test_zero_treated_as_positive(self):
|
||||
embs = [[0.0] * 8]
|
||||
result = build_typed_embeddings(embs, ["binary"])
|
||||
assert result.binary is not None
|
||||
# 0.0 >= 0 is True, so bit=1 for all => 127 (signed)
|
||||
assert result.binary[0] == [127]
|
||||
|
||||
def test_non_multiple_of_8_raises(self):
|
||||
embs = [[0.1] * 7]
|
||||
with pytest.raises(ValueError, match="multiple of 8"):
|
||||
build_typed_embeddings(embs, ["binary"])
|
||||
|
||||
def test_ubinary_non_multiple_of_8_raises(self):
|
||||
embs = [[0.1] * 10]
|
||||
with pytest.raises(ValueError, match="multiple of 8"):
|
||||
build_typed_embeddings(embs, ["ubinary"])
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsBase64:
|
||||
def test_base64_roundtrip(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["base64"])
|
||||
assert result.base64 is not None
|
||||
assert len(result.base64) == 2
|
||||
|
||||
for orig, b64_str in zip(sample_embeddings, result.base64):
|
||||
decoded = base64.b64decode(b64_str)
|
||||
n = len(orig)
|
||||
values = struct.unpack(f"<{n}f", decoded)
|
||||
np.testing.assert_allclose(orig, values, rtol=1e-5)
|
||||
|
||||
def test_base64_byte_length(self):
|
||||
embs = [[0.1, 0.2, 0.3]]
|
||||
result = build_typed_embeddings(embs, ["base64"])
|
||||
assert result.base64 is not None
|
||||
raw = base64.b64decode(result.base64[0])
|
||||
assert len(raw) == 3 * 4 # 3 floats * 4 bytes each
|
||||
|
||||
|
||||
class TestBuildTypedEmbeddingsMultiple:
|
||||
def test_all_types_at_once(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(
|
||||
sample_embeddings,
|
||||
["float", "binary", "ubinary", "base64"],
|
||||
)
|
||||
assert result.float is not None
|
||||
assert result.binary is not None
|
||||
assert result.ubinary is not None
|
||||
assert result.base64 is not None
|
||||
|
||||
def test_subset_types(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["float", "binary"])
|
||||
assert result.float is not None
|
||||
assert result.binary is not None
|
||||
assert result.ubinary is None
|
||||
assert result.base64 is None
|
||||
|
||||
def test_unknown_type_ignored(self, sample_embeddings: list[list[float]]):
|
||||
result = build_typed_embeddings(sample_embeddings, ["float", "unknown_type"])
|
||||
assert result.float is not None
|
||||
@@ -0,0 +1,9 @@
|
||||
model_name: "Qwen/Qwen3.5-35B-A3B"
|
||||
accuracy_threshold: 0.84
|
||||
tolerance: 0.03
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: >-
|
||||
--max-model-len 4096
|
||||
--data-parallel-size 2
|
||||
--enable-expert-parallel
|
||||
@@ -0,0 +1,10 @@
|
||||
model_name: "Qwen/Qwen3.5-35B-A3B-FP8"
|
||||
accuracy_threshold: 0.79
|
||||
tolerance: 0.03
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: >-
|
||||
--max-model-len 4096
|
||||
--data-parallel-size 2
|
||||
--enable-expert-parallel
|
||||
--kv-cache-dtype fp8
|
||||
@@ -0,0 +1,9 @@
|
||||
model_name: "nvidia/Qwen3.5-397B-A17B-NVFP4"
|
||||
accuracy_threshold: 0.88
|
||||
tolerance: 0.03
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: >-
|
||||
--max-model-len 4096
|
||||
--data-parallel-size 2
|
||||
--enable-expert-parallel
|
||||
@@ -0,0 +1,3 @@
|
||||
Qwen3.5-35B-A3B-DEP2.yaml
|
||||
Qwen3.5-35B-A3B-FP8-DEP2.yaml
|
||||
Qwen3.5-397B-A17B-NVFP4-DEP2.yaml
|
||||
@@ -19,8 +19,6 @@ from vllm.platforms import current_platform
|
||||
|
||||
from .gsm8k_eval import evaluate_gsm8k
|
||||
|
||||
TOL = 0.08 # Absolute tolerance for accuracy comparison
|
||||
|
||||
|
||||
def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict:
|
||||
"""Run GSM8K evaluation using our isolated script."""
|
||||
@@ -99,20 +97,20 @@ def test_gsm8k_correctness(config_filename):
|
||||
|
||||
measured_metric = results["accuracy"]
|
||||
expected_metric = eval_config["accuracy_threshold"]
|
||||
tol = eval_config.get("tolerance", 0.08)
|
||||
|
||||
print(f"GSM8K Results for {eval_config['model_name']}:")
|
||||
print(f" Measured metric: {measured_metric:.4f}")
|
||||
print(f" Expected metric: {expected_metric:.4f}")
|
||||
print(f" Tolerance: {TOL:.4f}")
|
||||
print(f" Tolerance: {tol:.4f}")
|
||||
print(f" Questions: {results['num_questions']}")
|
||||
print(f" Invalid rate: {results['invalid_rate']:.3f}")
|
||||
print(f" Latency: {results['latency']:.1f}s")
|
||||
print(f" QPS: {results['questions_per_second']:.1f}")
|
||||
|
||||
# Verify metric is within tolerance
|
||||
assert measured_metric >= expected_metric - TOL, (
|
||||
assert measured_metric >= expected_metric - tol, (
|
||||
f"GSM8K metric too low: {measured_metric:.4f} < "
|
||||
f"{expected_metric:.4f} - {TOL:.4f} = {expected_metric - TOL:.4f}"
|
||||
f"{expected_metric:.4f} - {tol:.4f} = {expected_metric - tol:.4f}"
|
||||
)
|
||||
|
||||
print(f"✅ GSM8K test passed for {eval_config['model_name']}")
|
||||
|
||||
@@ -33,7 +33,10 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.deep_gemm import is_deep_gemm_supported
|
||||
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
|
||||
from vllm.utils.flashinfer import (
|
||||
has_flashinfer_cutlass_fused_moe,
|
||||
has_flashinfer_nvlink_one_sided,
|
||||
)
|
||||
from vllm.utils.import_utils import (
|
||||
has_aiter,
|
||||
has_deep_ep,
|
||||
@@ -234,15 +237,15 @@ if has_mori():
|
||||
)
|
||||
|
||||
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_a2a_prepare_finalize import ( # noqa: E501
|
||||
FlashInferA2APrepareAndFinalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
|
||||
FlashInferExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_two_sided_prepare_finalize import ( # noqa: E501
|
||||
FlashInferNVLinkTwoSidedPrepareAndFinalize,
|
||||
)
|
||||
|
||||
register_prepare_and_finalize(
|
||||
FlashInferA2APrepareAndFinalize,
|
||||
FlashInferNVLinkTwoSidedPrepareAndFinalize,
|
||||
standard_format,
|
||||
nvfp4_types + fp8_types,
|
||||
blocked_quantization_support=True,
|
||||
@@ -263,6 +266,36 @@ else:
|
||||
FlashInferCutlassMoEPrepareAndFinalize = None
|
||||
FlashInferExperts = None
|
||||
|
||||
if (
|
||||
has_flashinfer_nvlink_one_sided()
|
||||
and has_flashinfer_cutlass_fused_moe()
|
||||
and current_platform.has_device_capability(100)
|
||||
):
|
||||
from vllm.model_executor.layers.fused_moe.flashinfer_nvlink_one_sided_prepare_finalize import ( # noqa: E501
|
||||
FlashInferNVLinkOneSidedPrepareAndFinalize,
|
||||
)
|
||||
|
||||
register_prepare_and_finalize(
|
||||
FlashInferNVLinkOneSidedPrepareAndFinalize,
|
||||
standard_format,
|
||||
nvfp4_types,
|
||||
blocked_quantization_support=False,
|
||||
backend="flashinfer_nvlink_one_sided",
|
||||
supports_apply_weight_on_input=False,
|
||||
)
|
||||
|
||||
if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability(100):
|
||||
from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
|
||||
TrtLlmNvFp4ExpertsModular,
|
||||
)
|
||||
|
||||
register_experts(
|
||||
TrtLlmNvFp4ExpertsModular,
|
||||
standard_format,
|
||||
nvfp4_types,
|
||||
blocked_quantization_support=False,
|
||||
supports_expert_map=True,
|
||||
)
|
||||
|
||||
if has_aiter():
|
||||
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
|
||||
|
||||
@@ -159,6 +159,52 @@ def test_quantize_to_fp4(
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[(32, 4096), (128, 4096), (1, 64), (127, 1024), (256, 16384)],
|
||||
)
|
||||
@pytest.mark.parametrize("is_sf_swizzled_layout", [True, False])
|
||||
@torch.inference_mode()
|
||||
def test_python_util_matches_cpp_allocation(
|
||||
shape: tuple[int, int],
|
||||
is_sf_swizzled_layout: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Verify that the Python utility (create_fp4_output_tensors) allocates
|
||||
tensors with the same shapes and dtypes as the C++ functional variant
|
||||
(scaled_fp4_quant_func).
|
||||
"""
|
||||
from vllm._custom_ops import create_fp4_output_tensors
|
||||
|
||||
torch.set_default_device("cuda:0")
|
||||
m, n = shape
|
||||
input_tensor = torch.randn((m, n), dtype=torch.bfloat16)
|
||||
input_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda:0")
|
||||
|
||||
# C++ functional variant allocates internally
|
||||
cpp_out, cpp_scale = torch.ops._C.scaled_fp4_quant(
|
||||
input_tensor, input_scale, is_sf_swizzled_layout
|
||||
)
|
||||
|
||||
# Python utility
|
||||
py_out, py_scale = create_fp4_output_tensors(
|
||||
m, n, torch.device("cuda:0"), is_sf_swizzled_layout
|
||||
)
|
||||
|
||||
assert py_out.shape == cpp_out.shape, (
|
||||
f"Output shape mismatch: Python {py_out.shape} vs C++ {cpp_out.shape}"
|
||||
)
|
||||
assert py_out.dtype == cpp_out.dtype, (
|
||||
f"Output dtype mismatch: Python {py_out.dtype} vs C++ {cpp_out.dtype}"
|
||||
)
|
||||
assert py_scale.shape == cpp_scale.shape, (
|
||||
f"Scale shape mismatch: Python {py_scale.shape} vs C++ {cpp_scale.shape}"
|
||||
)
|
||||
assert py_scale.dtype == cpp_scale.dtype, (
|
||||
f"Scale dtype mismatch: Python {py_scale.dtype} vs C++ {cpp_scale.dtype}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pad_shape", PAD_SHAPES)
|
||||
@torch.inference_mode()
|
||||
def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:
|
||||
|
||||
@@ -19,6 +19,7 @@ from vllm.transformers_utils.runai_utils import (
|
||||
def test_is_runai_obj_uri():
|
||||
assert is_runai_obj_uri("gs://some-gcs-bucket/path")
|
||||
assert is_runai_obj_uri("s3://some-s3-bucket/path")
|
||||
assert is_runai_obj_uri("az://some-azure-container/path")
|
||||
assert not is_runai_obj_uri("nfs://some-nfs-path")
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,361 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for EP weight filtering during model loading."""
|
||||
|
||||
import glob
|
||||
import tempfile
|
||||
|
||||
import huggingface_hub.constants
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.model_loader.ep_weight_filter import (
|
||||
compute_local_expert_ids,
|
||||
parse_expert_id,
|
||||
should_skip_weight,
|
||||
)
|
||||
from vllm.model_executor.model_loader.weight_utils import (
|
||||
safetensors_weights_iterator,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for parse_expert_id
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestParseExpertId:
|
||||
def test_routed_expert(self):
|
||||
name = "model.layers.0.mlp.experts.42.gate_proj.weight"
|
||||
assert parse_expert_id(name) == 42
|
||||
|
||||
def test_large_expert_id(self):
|
||||
name = "model.layers.60.mlp.experts.383.down_proj.weight"
|
||||
assert parse_expert_id(name) == 383
|
||||
|
||||
def test_shared_expert(self):
|
||||
# Shared experts use a different naming convention in most models
|
||||
name = "model.layers.0.mlp.shared_experts.gate_proj.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_attention_weight(self):
|
||||
name = "model.layers.0.self_attn.q_proj.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_embedding(self):
|
||||
name = "model.embed_tokens.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_layernorm(self):
|
||||
name = "model.layers.0.input_layernorm.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_fused_3d_expert(self):
|
||||
# 3D fused-expert tensors (e.g. gpt-oss) have no numeric expert id.
|
||||
# They must NOT be filtered — slicing happens later in weight_loader.
|
||||
name = "model.layers.0.mlp.experts.gate_proj.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_fused_3d_expert_down_proj(self):
|
||||
name = "model.layers.10.mlp.experts.down_proj.weight"
|
||||
assert parse_expert_id(name) is None
|
||||
|
||||
def test_expert_scale(self):
|
||||
# NVFP4 quantized models have scale tensors for experts
|
||||
name = "model.layers.5.mlp.experts.100.gate_proj.weight_scale"
|
||||
assert parse_expert_id(name) == 100
|
||||
|
||||
def test_expert_zero_id(self):
|
||||
name = "model.layers.0.mlp.experts.0.up_proj.weight"
|
||||
assert parse_expert_id(name) == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for compute_local_expert_ids
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestComputeLocalExpertIds:
|
||||
def test_ep_disabled(self):
|
||||
assert compute_local_expert_ids(64, ep_size=1, ep_rank=0) is None
|
||||
|
||||
def test_even_split(self):
|
||||
# 64 experts, EP=8 → 8 per rank
|
||||
ids = compute_local_expert_ids(64, ep_size=8, ep_rank=0)
|
||||
assert ids == set(range(0, 8))
|
||||
|
||||
ids = compute_local_expert_ids(64, ep_size=8, ep_rank=7)
|
||||
assert ids == set(range(56, 64))
|
||||
|
||||
def test_uneven_split(self):
|
||||
# 10 experts, EP=3 → ranks get 4, 3, 3
|
||||
ids_0 = compute_local_expert_ids(10, ep_size=3, ep_rank=0)
|
||||
ids_1 = compute_local_expert_ids(10, ep_size=3, ep_rank=1)
|
||||
ids_2 = compute_local_expert_ids(10, ep_size=3, ep_rank=2)
|
||||
|
||||
assert len(ids_0) == 4
|
||||
assert len(ids_1) == 3
|
||||
assert len(ids_2) == 3
|
||||
# All experts covered, no overlap
|
||||
assert ids_0 | ids_1 | ids_2 == set(range(10))
|
||||
assert ids_0.isdisjoint(ids_1)
|
||||
assert ids_1.isdisjoint(ids_2)
|
||||
|
||||
def test_384_experts_ep8(self):
|
||||
# Kimi-K2.5 config: 384 experts, EP=8
|
||||
for rank in range(8):
|
||||
ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
|
||||
assert len(ids) == 48
|
||||
|
||||
# All experts covered
|
||||
all_ids = set()
|
||||
for rank in range(8):
|
||||
ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
|
||||
all_ids |= ids
|
||||
assert all_ids == set(range(384))
|
||||
|
||||
def test_384_experts_ep16(self):
|
||||
for rank in range(16):
|
||||
ids = compute_local_expert_ids(384, ep_size=16, ep_rank=rank)
|
||||
assert len(ids) == 24
|
||||
|
||||
def test_384_experts_ep24(self):
|
||||
# 384 / 24 = 16 exactly
|
||||
for rank in range(24):
|
||||
ids = compute_local_expert_ids(384, ep_size=24, ep_rank=rank)
|
||||
assert len(ids) == 16
|
||||
|
||||
# round_robin placement tests
|
||||
|
||||
def test_round_robin_basic(self):
|
||||
# 8 experts, EP=2: rank 0 → {0,2,4,6}, rank 1 → {1,3,5,7}
|
||||
rr = "round_robin"
|
||||
ids_0 = compute_local_expert_ids(8, 2, 0, placement=rr)
|
||||
ids_1 = compute_local_expert_ids(8, 2, 1, placement=rr)
|
||||
assert ids_0 == {0, 2, 4, 6}
|
||||
assert ids_1 == {1, 3, 5, 7}
|
||||
|
||||
def test_round_robin_full_coverage(self):
|
||||
# 384 experts, EP=8: all experts covered, no overlap
|
||||
rr = "round_robin"
|
||||
all_ids: set[int] = set()
|
||||
for rank in range(8):
|
||||
ids = compute_local_expert_ids(384, 8, rank, placement=rr)
|
||||
assert ids is not None and len(ids) == 48
|
||||
assert all_ids.isdisjoint(ids)
|
||||
all_ids |= ids
|
||||
assert all_ids == set(range(384))
|
||||
|
||||
def test_round_robin_uneven(self):
|
||||
# 10 experts, EP=3: rank 0→{0,3,6,9}, rank 1→{1,4,7}, rank 2→{2,5,8}
|
||||
rr = "round_robin"
|
||||
ids_0 = compute_local_expert_ids(10, 3, 0, placement=rr)
|
||||
ids_1 = compute_local_expert_ids(10, 3, 1, placement=rr)
|
||||
ids_2 = compute_local_expert_ids(10, 3, 2, placement=rr)
|
||||
assert ids_0 == {0, 3, 6, 9}
|
||||
assert ids_1 == {1, 4, 7}
|
||||
assert ids_2 == {2, 5, 8}
|
||||
assert ids_0 | ids_1 | ids_2 == set(range(10))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests for should_skip_weight
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestShouldSkipWeight:
|
||||
def setup_method(self):
|
||||
# Simulate EP=8, rank=0 → experts 0-47
|
||||
self.local_ids = compute_local_expert_ids(384, ep_size=8, ep_rank=0)
|
||||
|
||||
def test_no_filter(self):
|
||||
assert not should_skip_weight("anything", None)
|
||||
|
||||
def test_dense_not_skipped(self):
|
||||
assert not should_skip_weight(
|
||||
"model.layers.0.self_attn.q_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
def test_local_expert_not_skipped(self):
|
||||
assert not should_skip_weight(
|
||||
"model.layers.0.mlp.experts.10.gate_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
def test_remote_expert_skipped(self):
|
||||
assert should_skip_weight(
|
||||
"model.layers.0.mlp.experts.200.gate_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
def test_boundary_expert(self):
|
||||
# Expert 47 is local (last one), 48 is not
|
||||
assert not should_skip_weight(
|
||||
"model.layers.0.mlp.experts.47.gate_proj.weight", self.local_ids
|
||||
)
|
||||
assert should_skip_weight(
|
||||
"model.layers.0.mlp.experts.48.gate_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
def test_shared_expert_not_skipped(self):
|
||||
assert not should_skip_weight(
|
||||
"model.layers.0.mlp.shared_experts.gate_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
def test_embedding_not_skipped(self):
|
||||
assert not should_skip_weight("model.embed_tokens.weight", self.local_ids)
|
||||
|
||||
def test_fused_3d_expert_not_skipped(self):
|
||||
# 3D fused-expert tensors (gpt-oss style) have no numeric id.
|
||||
# Must not be skipped — weight_loader handles slicing later.
|
||||
assert not should_skip_weight(
|
||||
"model.layers.0.mlp.experts.gate_proj.weight", self.local_ids
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration test: safetensors_weights_iterator with EP filtering
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSafetensorsWeightsIteratorWithEpFilter:
|
||||
"""Verify that EP filtering produces a strict subset of unfiltered loading
|
||||
and that all expected dense + local expert weights are present."""
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def gpt2_files(self):
|
||||
"""Download GPT-2 safetensors to a temp dir (shared across class)."""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
huggingface_hub.constants.HF_HUB_OFFLINE = False
|
||||
from vllm.model_executor.model_loader.weight_utils import (
|
||||
download_weights_from_hf,
|
||||
)
|
||||
|
||||
download_weights_from_hf(
|
||||
"openai-community/gpt2",
|
||||
allow_patterns=["*.safetensors"],
|
||||
cache_dir=tmpdir,
|
||||
)
|
||||
files = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
|
||||
assert len(files) > 0
|
||||
yield files
|
||||
|
||||
def test_no_filter_returns_all(self, gpt2_files):
|
||||
"""With local_expert_ids=None, all weights are returned (no MoE)."""
|
||||
all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
|
||||
filtered_weights = dict(
|
||||
safetensors_weights_iterator(gpt2_files, False, local_expert_ids=None)
|
||||
)
|
||||
assert set(all_weights.keys()) == set(filtered_weights.keys())
|
||||
|
||||
def test_empty_filter_skips_experts_only(self, gpt2_files):
|
||||
"""GPT-2 has no expert weights, so even an empty local_expert_ids
|
||||
set should return all weights (all are dense)."""
|
||||
all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
|
||||
filtered_weights = dict(
|
||||
safetensors_weights_iterator(gpt2_files, False, local_expert_ids=set())
|
||||
)
|
||||
# GPT-2 has no experts, so nothing should be filtered
|
||||
assert set(all_weights.keys()) == set(filtered_weights.keys())
|
||||
|
||||
|
||||
class TestEpFilterOnSyntheticMoeWeights:
|
||||
"""Create synthetic safetensors files with expert-like naming and verify
|
||||
that the filter correctly skips non-local experts."""
|
||||
|
||||
@pytest.fixture
|
||||
def synthetic_moe_files(self, tmp_path):
|
||||
"""Create synthetic safetensors with expert-patterned tensor names."""
|
||||
from safetensors.torch import save_file
|
||||
|
||||
tensors = {}
|
||||
# Dense weights
|
||||
tensors["model.embed_tokens.weight"] = torch.randn(100, 64)
|
||||
tensors["model.layers.0.self_attn.q_proj.weight"] = torch.randn(64, 64)
|
||||
tensors["model.layers.0.input_layernorm.weight"] = torch.randn(64)
|
||||
# Expert weights: 8 experts
|
||||
for expert_id in range(8):
|
||||
tensors[f"model.layers.0.mlp.experts.{expert_id}.gate_proj.weight"] = (
|
||||
torch.randn(128, 64)
|
||||
)
|
||||
tensors[f"model.layers.0.mlp.experts.{expert_id}.up_proj.weight"] = (
|
||||
torch.randn(128, 64)
|
||||
)
|
||||
tensors[f"model.layers.0.mlp.experts.{expert_id}.down_proj.weight"] = (
|
||||
torch.randn(64, 128)
|
||||
)
|
||||
# Shared expert (should never be filtered)
|
||||
tensors["model.layers.0.mlp.shared_experts.gate_proj.weight"] = torch.randn(
|
||||
128, 64
|
||||
)
|
||||
|
||||
filepath = str(tmp_path / "model-00001-of-00001.safetensors")
|
||||
save_file(tensors, filepath)
|
||||
return [filepath], tensors
|
||||
|
||||
def test_no_filter_returns_all(self, synthetic_moe_files):
|
||||
files, expected = synthetic_moe_files
|
||||
loaded = dict(safetensors_weights_iterator(files, False))
|
||||
assert set(loaded.keys()) == set(expected.keys())
|
||||
|
||||
def test_ep2_rank0_gets_half_experts(self, synthetic_moe_files):
|
||||
files, expected = synthetic_moe_files
|
||||
# EP=2, rank=0 → experts 0-3
|
||||
local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
|
||||
loaded = dict(
|
||||
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
|
||||
)
|
||||
|
||||
# Should have all dense + shared + experts 0-3 only
|
||||
for name in loaded:
|
||||
eid = parse_expert_id(name)
|
||||
if eid is not None:
|
||||
assert eid in local_ids, f"Non-local expert {eid} was loaded"
|
||||
|
||||
# Check expert count: 4 experts × 3 weights = 12
|
||||
expert_names = [n for n in loaded if parse_expert_id(n) is not None]
|
||||
assert len(expert_names) == 4 * 3
|
||||
|
||||
# Check all dense weights present
|
||||
assert "model.embed_tokens.weight" in loaded
|
||||
assert "model.layers.0.self_attn.q_proj.weight" in loaded
|
||||
assert "model.layers.0.input_layernorm.weight" in loaded
|
||||
assert "model.layers.0.mlp.shared_experts.gate_proj.weight" in loaded
|
||||
|
||||
def test_ep2_rank1_gets_other_half(self, synthetic_moe_files):
|
||||
files, expected = synthetic_moe_files
|
||||
local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=1)
|
||||
loaded = dict(
|
||||
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
|
||||
)
|
||||
|
||||
expert_names = [n for n in loaded if parse_expert_id(n) is not None]
|
||||
assert len(expert_names) == 4 * 3
|
||||
for name in expert_names:
|
||||
assert parse_expert_id(name) in local_ids
|
||||
|
||||
def test_ep8_each_rank_gets_one_expert(self, synthetic_moe_files):
|
||||
files, _ = synthetic_moe_files
|
||||
all_expert_names = set()
|
||||
for rank in range(8):
|
||||
local_ids = compute_local_expert_ids(8, ep_size=8, ep_rank=rank)
|
||||
loaded = dict(
|
||||
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
|
||||
)
|
||||
expert_names = {n for n in loaded if parse_expert_id(n) is not None}
|
||||
# 1 expert × 3 weights
|
||||
assert len(expert_names) == 3
|
||||
all_expert_names |= expert_names
|
||||
|
||||
# All 8 experts × 3 weights covered across ranks
|
||||
assert len(all_expert_names) == 24
|
||||
|
||||
def test_tensor_values_match(self, synthetic_moe_files):
|
||||
"""Filtered tensors have identical values to unfiltered ones."""
|
||||
files, _ = synthetic_moe_files
|
||||
all_weights = dict(safetensors_weights_iterator(files, False))
|
||||
|
||||
local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
|
||||
filtered = dict(
|
||||
safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
|
||||
)
|
||||
|
||||
for name, tensor in filtered.items():
|
||||
assert torch.equal(tensor, all_weights[name]), f"Tensor mismatch for {name}"
|
||||
@@ -0,0 +1,68 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for CPU unquantized GEMM dispatch behavior."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers import utils
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def _mock_zentorch_linear_unary():
|
||||
"""Register a mock zentorch_linear_unary op when zentorch is not installed.
|
||||
|
||||
Allows the dispatch tests to run in CI without a real zentorch build.
|
||||
Skips registration when zentorch is already available.
|
||||
"""
|
||||
if hasattr(torch.ops.zentorch, "zentorch_linear_unary"):
|
||||
yield
|
||||
return
|
||||
|
||||
lib_def = torch.library.Library("zentorch", "DEF")
|
||||
lib_def.define(
|
||||
"zentorch_linear_unary("
|
||||
"Tensor input, "
|
||||
"Tensor weight, "
|
||||
"Tensor? bias, "
|
||||
"bool is_weight_prepacked=False"
|
||||
") -> Tensor"
|
||||
)
|
||||
|
||||
lib_impl = torch.library.Library("zentorch", "IMPL", "CPU")
|
||||
lib_impl.impl(
|
||||
"zentorch_linear_unary",
|
||||
lambda input, weight, bias, is_weight_prepacked=False: (
|
||||
torch.nn.functional.linear(input, weight, bias)
|
||||
),
|
||||
)
|
||||
|
||||
yield
|
||||
|
||||
lib_impl._destroy()
|
||||
lib_def._destroy()
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
|
||||
def test_dispatch_cpu_unquantized_gemm_uses_zentorch_on_zen(monkeypatch):
|
||||
monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
|
||||
|
||||
layer = torch.nn.Linear(16, 8, bias=True)
|
||||
x = torch.randn(4, 16)
|
||||
expected = torch.nn.functional.linear(x, layer.weight, layer.bias)
|
||||
|
||||
utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
|
||||
output = layer.cpu_linear(x, layer.weight, layer.bias)
|
||||
|
||||
torch.testing.assert_close(output, expected)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
|
||||
def test_dispatch_cpu_unquantized_gemm_zen_remove_weight(monkeypatch):
|
||||
monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
|
||||
|
||||
layer = torch.nn.Linear(16, 8, bias=True)
|
||||
utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
|
||||
|
||||
assert layer.weight.numel() == 0
|
||||
@@ -103,6 +103,10 @@ AITER_MODEL_LIST = [
|
||||
marks=[pytest.mark.core_model, pytest.mark.cpu_model],
|
||||
),
|
||||
pytest.param("swiss-ai/Apertus-8B-Instruct-2509"), # apertus
|
||||
pytest.param(
|
||||
"naver-hyperclovax/HyperCLOVAX-SEED-Think-14B", # hyperclovax
|
||||
marks=[large_gpu_mark(min_gb=32)],
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("max_tokens", [32])
|
||||
|
||||
@@ -777,6 +777,7 @@ VLM_TEST_SETTINGS = {
|
||||
max_model_len=8192,
|
||||
max_num_seqs=2,
|
||||
auto_cls=AutoModelForCausalLM,
|
||||
patch_hf_runner=model_utils.paddleocr_vl_patch_hf_runner,
|
||||
image_size_factors=[(0.25,)],
|
||||
marks=[
|
||||
pytest.mark.skipif(
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import contextlib
|
||||
from dataclasses import asdict
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from mistral_common.audio import Audio
|
||||
from mistral_common.protocol.instruct.chunk import RawAudio
|
||||
from mistral_common.protocol.transcription.request import (
|
||||
@@ -17,18 +19,21 @@ from vllm.assets.audio import AudioAsset
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
|
||||
from ....utils import ROCM_ENGINE_KWARGS
|
||||
|
||||
MODEL_NAME = "mistralai/Voxtral-Mini-4B-Realtime-2602"
|
||||
ENGINE_CONFIG = dict(
|
||||
model=MODEL_NAME,
|
||||
max_model_len=8192,
|
||||
max_num_seqs=4,
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
config_format="mistral",
|
||||
load_format="mistral",
|
||||
tokenizer_mode="mistral",
|
||||
enforce_eager=True,
|
||||
gpu_memory_utilization=0.9,
|
||||
)
|
||||
ENGINE_CONFIG = {
|
||||
"model": MODEL_NAME,
|
||||
"max_model_len": 8192,
|
||||
"max_num_seqs": 4,
|
||||
"limit_mm_per_prompt": {"audio": 1},
|
||||
"config_format": "mistral",
|
||||
"load_format": "mistral",
|
||||
"tokenizer_mode": "mistral",
|
||||
"enforce_eager": True,
|
||||
"gpu_memory_utilization": 0.9,
|
||||
**ROCM_ENGINE_KWARGS,
|
||||
}
|
||||
|
||||
|
||||
EXPECTED_TEXT = [
|
||||
@@ -49,6 +54,14 @@ EXPECTED_TEXT = [
|
||||
]
|
||||
|
||||
|
||||
def _normalize(texts: list[str]) -> list[str]:
|
||||
# The model occasionally transcribes "OBS" as "a base hit" and
|
||||
# "oh, my" as "oh my", but both are acoustically valid. Normalise so
|
||||
# the assertion is stable across runs and hardware.
|
||||
texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my")
|
||||
return texts
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def audio_assets() -> list[AudioAsset]:
|
||||
return [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
|
||||
@@ -60,15 +73,27 @@ def tokenizer() -> MistralTokenizer:
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def engine() -> LLM:
|
||||
def engine():
|
||||
engine_args = EngineArgs(**ENGINE_CONFIG)
|
||||
return LLM(**asdict(engine_args))
|
||||
llm = LLM(**asdict(engine_args))
|
||||
try:
|
||||
yield llm
|
||||
finally:
|
||||
with contextlib.suppress(Exception):
|
||||
llm.llm_engine.engine_core.shutdown()
|
||||
import torch
|
||||
|
||||
torch.accelerator.empty_cache()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_engine() -> AsyncLLM:
|
||||
@pytest_asyncio.fixture
|
||||
async def async_engine():
|
||||
engine_args = AsyncEngineArgs(**ENGINE_CONFIG)
|
||||
return AsyncLLM.from_engine_args(engine_args)
|
||||
llm = AsyncLLM.from_engine_args(engine_args)
|
||||
try:
|
||||
yield llm
|
||||
finally:
|
||||
llm.shutdown()
|
||||
|
||||
|
||||
def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
|
||||
@@ -108,8 +133,13 @@ def test_voxtral_realtime_forward(audio_assets, tokenizer, engine):
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
|
||||
texts = [out.outputs[0].text for out in outputs]
|
||||
assert texts == EXPECTED_TEXT
|
||||
texts = _normalize([out.outputs[0].text for out in outputs])
|
||||
for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
|
||||
assert got == expected, (
|
||||
f"Output mismatch at index {i}:\n"
|
||||
f" got: {got!r}\n"
|
||||
f" expected: {expected!r}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -149,9 +179,17 @@ async def test_voxtral_realtime_generator(audio_assets, tokenizer, async_engine)
|
||||
|
||||
output_tokens_list.append(output_tokens)
|
||||
|
||||
texts = [
|
||||
tokenizer.decode(output_tokens, special_token_policy=SpecialTokenPolicy.IGNORE)
|
||||
for output_tokens in output_tokens_list
|
||||
]
|
||||
texts[1] = texts[1].replace("a base hit", "OBS").replace("oh my", "oh, my")
|
||||
assert texts == EXPECTED_TEXT
|
||||
texts = _normalize(
|
||||
[
|
||||
tokenizer.decode(
|
||||
output_tokens, special_token_policy=SpecialTokenPolicy.IGNORE
|
||||
)
|
||||
for output_tokens in output_tokens_list
|
||||
]
|
||||
)
|
||||
for i, (got, expected) in enumerate(zip(texts, EXPECTED_TEXT)):
|
||||
assert got == expected, (
|
||||
f"Output mismatch at index {i}:\n"
|
||||
f" got: {got!r}\n"
|
||||
f" expected: {expected!r}"
|
||||
)
|
||||
|
||||
@@ -24,6 +24,7 @@ from transformers import (
|
||||
GenerationConfig,
|
||||
GenerationMixin,
|
||||
)
|
||||
from transformers.masking_utils import create_causal_mask
|
||||
from transformers.video_utils import VideoMetadata
|
||||
|
||||
from vllm.logprobs import SampleLogprobs
|
||||
@@ -679,10 +680,14 @@ def isaac_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
|
||||
sin = sin.to(inputs_embeds.dtype)
|
||||
|
||||
# Prepare attention mask
|
||||
if attention_mask is not None:
|
||||
attention_mask = self._update_causal_mask(
|
||||
attention_mask, inputs_embeds, cache_position, past_key_values, False
|
||||
)
|
||||
attention_mask = create_causal_mask(
|
||||
config=self.config,
|
||||
input_embeds=inputs_embeds,
|
||||
attention_mask=attention_mask,
|
||||
past_key_values=past_key_values,
|
||||
position_ids=position_ids,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
# Initialize and collect hidden states
|
||||
hidden_states = inputs_embeds
|
||||
@@ -1149,6 +1154,31 @@ def ovis2_5_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
|
||||
return hf_model
|
||||
|
||||
|
||||
def paddleocr_vl_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
|
||||
"""Patches the HfRunner to fix create_causal_mask API mismatch.
|
||||
|
||||
The PaddleOCR-VL HF model passes `inputs_embeds` to create_causal_mask,
|
||||
but transformers renamed this parameter to `input_embeds`.
|
||||
"""
|
||||
import sys
|
||||
|
||||
model_module = sys.modules.get(type(hf_model.model.model).__module__)
|
||||
if model_module is None:
|
||||
return hf_model
|
||||
|
||||
original_create_causal_mask = getattr(model_module, "create_causal_mask", None)
|
||||
if original_create_causal_mask is None:
|
||||
return hf_model
|
||||
|
||||
def patched_create_causal_mask(*args, **kwargs):
|
||||
if "inputs_embeds" in kwargs:
|
||||
kwargs["input_embeds"] = kwargs.pop("inputs_embeds")
|
||||
return original_create_causal_mask(*args, **kwargs)
|
||||
|
||||
model_module.create_causal_mask = patched_create_causal_mask # type: ignore[attr-defined]
|
||||
return hf_model
|
||||
|
||||
|
||||
def qwen2_5_omni_patch_hf_runner(hf_model: HfRunner) -> HfRunner:
|
||||
"""Patches and returns an instance of the HfRunner for Qwen2.5-Omni."""
|
||||
thinker = hf_model.model.thinker
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import CLIPModel
|
||||
|
||||
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
|
||||
@@ -50,13 +51,16 @@ def _run_test(
|
||||
if "pixel_values" in inputs:
|
||||
pooled_output = hf_model.model.get_image_features(
|
||||
pixel_values=inputs.pixel_values,
|
||||
).squeeze(0)
|
||||
)
|
||||
else:
|
||||
pooled_output = hf_model.model.get_text_features(
|
||||
input_ids=inputs.input_ids,
|
||||
attention_mask=inputs.attention_mask,
|
||||
).squeeze(0)
|
||||
)
|
||||
|
||||
if not isinstance(pooled_output, torch.Tensor):
|
||||
pooled_output = pooled_output.pooler_output
|
||||
pooled_output = pooled_output.squeeze(0)
|
||||
all_outputs.append(pooled_output.tolist())
|
||||
|
||||
hf_outputs = all_outputs
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import SiglipModel
|
||||
|
||||
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
|
||||
@@ -68,12 +69,15 @@ def _run_test(
|
||||
if "pixel_values" in inputs:
|
||||
pooled_output = hf_model.model.get_image_features(
|
||||
pixel_values=inputs.pixel_values,
|
||||
).squeeze(0)
|
||||
)
|
||||
else:
|
||||
pooled_output = hf_model.model.get_text_features(
|
||||
input_ids=inputs.input_ids,
|
||||
).squeeze(0)
|
||||
)
|
||||
|
||||
if not isinstance(pooled_output, torch.Tensor):
|
||||
pooled_output = pooled_output.pooler_output
|
||||
pooled_output = pooled_output.squeeze(0)
|
||||
all_outputs.append(pooled_output.tolist())
|
||||
|
||||
hf_outputs = all_outputs
|
||||
|
||||
@@ -34,8 +34,22 @@ MODELS = [
|
||||
]
|
||||
|
||||
|
||||
def create_mm_data(num_videos: int) -> dict[str, list]:
|
||||
# Small video (8 frames, 64×64) and ~0.5 s of audio at 16 kHz so the test
|
||||
# stays fast even without a GPU.
|
||||
mm_data = dict[str, list](video=[], audio=[])
|
||||
for i in range(num_videos):
|
||||
rng = np.random.RandomState(i)
|
||||
video = random_video(rng, min_frames=8, max_frames=9, min_wh=64, max_wh=65)
|
||||
audio, sr = random_audio(rng, min_len=8000, max_len=8001, sr=16000)
|
||||
mm_data["video"].append(video)
|
||||
mm_data["audio"].append((audio, sr))
|
||||
return mm_data
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_id", MODELS)
|
||||
def test_audio_in_video_cache_correctness(model_id: str) -> None:
|
||||
@pytest.mark.parametrize("num_videos", [1, 2])
|
||||
def test_audio_in_video_cache_correctness(model_id: str, num_videos: int) -> None:
|
||||
"""
|
||||
Regression test for https://github.com/vllm-project/vllm/pull/36800
|
||||
|
||||
@@ -47,7 +61,7 @@ def test_audio_in_video_cache_correctness(model_id: str) -> None:
|
||||
"""
|
||||
ctx = build_model_context(
|
||||
model_id,
|
||||
limit_mm_per_prompt={"audio": 1, "image": 0, "video": 1},
|
||||
limit_mm_per_prompt={"audio": num_videos, "image": 0, "video": num_videos},
|
||||
mm_processor_cache_gb=1,
|
||||
)
|
||||
|
||||
@@ -65,17 +79,12 @@ def test_audio_in_video_cache_correctness(model_id: str) -> None:
|
||||
|
||||
video_token_id = baseline_processor.info.get_hf_config().video_token_id
|
||||
|
||||
rng = np.random.RandomState(0)
|
||||
# Small video (8 frames, 64×64) and ~0.5 s of audio at 16 kHz so the test
|
||||
# stays fast even without a GPU.
|
||||
video = random_video(rng, min_frames=8, max_frames=9, min_wh=64, max_wh=65)
|
||||
audio, sr = random_audio(rng, min_len=8000, max_len=8001, sr=16000)
|
||||
mm_data = {"video": [video], "audio": [(audio, sr)]}
|
||||
mm_data = create_mm_data(num_videos)
|
||||
hf_processor_mm_kwargs = {"use_audio_in_video": True}
|
||||
|
||||
def run(processor):
|
||||
return processor(
|
||||
[video_token_id],
|
||||
[video_token_id] * num_videos,
|
||||
mm_items=baseline_processor.info.parse_mm_data(mm_data),
|
||||
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
|
||||
)["prompt_token_ids"]
|
||||
|
||||
@@ -320,7 +320,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
|
||||
"tencent/Hunyuan-A13B-Instruct", trust_remote_code=True
|
||||
),
|
||||
"HyperCLOVAXForCausalLM": _HfExamplesInfo(
|
||||
"naver-hyperclovax/HyperCLOVAX-SEED-Think-32B",
|
||||
"naver-hyperclovax/HyperCLOVAX-SEED-Think-14B",
|
||||
trust_remote_code=True,
|
||||
),
|
||||
"InternLMForCausalLM": _HfExamplesInfo(
|
||||
@@ -780,6 +780,7 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
"Ernie4_5_VLMoeForConditionalGeneration": _HfExamplesInfo(
|
||||
"baidu/ERNIE-4.5-VL-28B-A3B-PT",
|
||||
trust_remote_code=True,
|
||||
revision="refs/pr/17",
|
||||
),
|
||||
"FireRedASR2ForConditionalGeneration": _HfExamplesInfo(
|
||||
"allendou/FireRedASR2-LLM-vllm",
|
||||
|
||||
@@ -4,6 +4,7 @@ import base64
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
@@ -71,3 +72,13 @@ def test_audio_media_io_encode_base64(dummy_audio):
|
||||
decoded = base64.b64decode(out)
|
||||
assert decoded == b"dummy_wav_data"
|
||||
mock_write.assert_called_once()
|
||||
|
||||
|
||||
def test_audio_media_io_from_video(video_assets):
|
||||
audio_io = AudioMediaIO()
|
||||
video_path = video_assets[0].video_path
|
||||
with open(video_path, "rb") as f:
|
||||
audio, sr = audio_io.load_bytes(f.read())
|
||||
audio_ref, sr_ref = librosa.load(video_path, sr=None)
|
||||
assert sr == sr_ref
|
||||
np.testing.assert_allclose(audio_ref, audio, atol=1e-4)
|
||||
|
||||
@@ -1,11 +1,19 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import json
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm.entrypoints.mcp.tool_server import ToolServer
|
||||
from vllm.reasoning import ReasoningParser
|
||||
from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser
|
||||
from vllm.reasoning.gptoss_reasoning_parser import (
|
||||
GptOssReasoningParser,
|
||||
from_builtin_tool_to_tag,
|
||||
no_func_reasoning_tag,
|
||||
)
|
||||
|
||||
REASONING_MODEL_NAME = "openai/gpt-oss-120b"
|
||||
|
||||
@@ -142,3 +150,133 @@ def test_gptoss_is_reasoning_end(
|
||||
output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(output)
|
||||
actual_is_reasoning_end = parser.is_reasoning_end(output_ids)
|
||||
assert is_reasoning_end == actual_is_reasoning_end
|
||||
|
||||
|
||||
class TestGptOssStructuralTags:
|
||||
"""Test cases for GptOssReasoningParser structural tag functionality."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tokenizer(self):
|
||||
"""Create a mock tokenizer for testing."""
|
||||
tokenizer = Mock()
|
||||
tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5])
|
||||
tokenizer.get_vocab = Mock(return_value={"<|end|>": 6})
|
||||
return tokenizer
|
||||
|
||||
@pytest.fixture
|
||||
def reasoning_parser(self, mock_tokenizer):
|
||||
"""Create a GptOssReasoningParser instance."""
|
||||
return GptOssReasoningParser(mock_tokenizer)
|
||||
|
||||
def test_prepare_structured_tag_no_tool_server(self, reasoning_parser):
|
||||
"""Test prepare_structured_tag with no tool server."""
|
||||
result = reasoning_parser.prepare_structured_tag(None, None)
|
||||
expected = json.dumps(no_func_reasoning_tag)
|
||||
|
||||
assert result == expected
|
||||
|
||||
# Verify the structure is correct
|
||||
parsed = json.loads(result)
|
||||
assert parsed["type"] == "structural_tag"
|
||||
assert parsed["format"]["type"] == "triggered_tags"
|
||||
assert len(parsed["format"]["tags"]) == 1
|
||||
assert parsed["format"]["tags"][0]["begin"] == "<|channel|>analysis<|message|>"
|
||||
assert parsed["format"]["triggers"] == ["<|channel|>analysis"]
|
||||
|
||||
def test_prepare_structured_tag_with_original_tag(self, reasoning_parser):
|
||||
"""Test prepare_structured_tag when original_tag is provided."""
|
||||
original_tag = '{"custom": "tag"}'
|
||||
result = reasoning_parser.prepare_structured_tag(original_tag, None)
|
||||
|
||||
# Should return the original tag unchanged
|
||||
assert result == original_tag
|
||||
|
||||
def test_from_builtin_tool_to_tag(self):
|
||||
"""Test from_builtin_tool_to_tag function."""
|
||||
tags = from_builtin_tool_to_tag("python")
|
||||
|
||||
assert len(tags) == 2
|
||||
assert tags[0]["begin"] == "<|channel|>commentary to=python"
|
||||
assert tags[0]["content"]["type"] == "any_text"
|
||||
assert tags[0]["end"] == "<|end|>"
|
||||
|
||||
assert tags[1]["begin"] == "<|channel|>analysis to=python"
|
||||
assert tags[1]["content"]["type"] == "any_text"
|
||||
assert tags[1]["end"] == "<|end|>"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tools",
|
||||
[
|
||||
[],
|
||||
["browser"],
|
||||
["python"],
|
||||
["container"],
|
||||
["browser", "python"],
|
||||
["browser", "container"],
|
||||
["python", "container"],
|
||||
["browser", "python", "container"],
|
||||
],
|
||||
)
|
||||
def test_json_validity_comprehensive(self, reasoning_parser, tools):
|
||||
"""Test JSON validity across all possible tool combinations."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=lambda tool: tool in tools)
|
||||
|
||||
result = reasoning_parser.prepare_structured_tag(None, tool_server)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
assert parsed_result["type"] == "structural_tag"
|
||||
assert "format" in parsed_result
|
||||
assert "tags" in parsed_result["format"]
|
||||
assert "triggers" in parsed_result["format"]
|
||||
|
||||
# Tag count should be: 1 (analysis) + 2 * len(tools)
|
||||
expected_tag_count = 1 + (2 * len(tools))
|
||||
assert len(parsed_result["format"]["tags"]) == expected_tag_count
|
||||
|
||||
# Verify triggers are correctly configured
|
||||
expected_triggers = ["<|channel|>analysis"]
|
||||
if tools:
|
||||
expected_triggers.append("<|channel|>commentary to=")
|
||||
assert set(parsed_result["format"]["triggers"]) == set(expected_triggers)
|
||||
|
||||
def test_no_cross_request_state_pollution(self, reasoning_parser):
|
||||
"""Test that sequential calls with different tool servers produce
|
||||
independent results, guarding against shared mutable state
|
||||
(e.g. missing deepcopy in tag_with_builtin_funcs)."""
|
||||
tool_server_1 = Mock(spec=ToolServer)
|
||||
tool_server_1.has_tool = Mock(side_effect=lambda tool: tool == "python")
|
||||
|
||||
tool_server_2 = Mock(spec=ToolServer)
|
||||
tool_server_2.has_tool = Mock(side_effect=lambda tool: tool == "browser")
|
||||
|
||||
result_1 = reasoning_parser.prepare_structured_tag(None, tool_server_1)
|
||||
result_2 = reasoning_parser.prepare_structured_tag(None, tool_server_2)
|
||||
|
||||
tags_1 = [tag["begin"] for tag in json.loads(result_1)["format"]["tags"]]
|
||||
tags_2 = [tag["begin"] for tag in json.loads(result_2)["format"]["tags"]]
|
||||
|
||||
assert "<|channel|>commentary to=python" in tags_1
|
||||
assert "<|channel|>commentary to=browser" not in tags_1
|
||||
|
||||
assert "<|channel|>commentary to=browser" in tags_2
|
||||
assert "<|channel|>commentary to=python" not in tags_2
|
||||
|
||||
def test_tag_format_consistency(self, reasoning_parser):
|
||||
"""Test that all generated tags follow consistent format,
|
||||
catching malformed tags from from_builtin_tool_to_tag."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(
|
||||
side_effect=lambda tool: tool in ["python", "browser"]
|
||||
)
|
||||
|
||||
result = reasoning_parser.prepare_structured_tag(None, tool_server)
|
||||
parsed_result = json.loads(result)
|
||||
|
||||
for tag in parsed_result["format"]["tags"]:
|
||||
assert "begin" in tag
|
||||
assert "content" in tag
|
||||
assert "end" in tag
|
||||
assert tag["content"]["type"] == "any_text"
|
||||
assert tag["end"] == "<|end|>"
|
||||
assert tag["begin"].startswith("<|channel|>")
|
||||
|
||||
@@ -55,7 +55,7 @@ def test_gc():
|
||||
# The memory allocated for model and KV cache should be released.
|
||||
# The memory allocated for PyTorch and others should be less than 50MB.
|
||||
# Usually, it's around 10MB.
|
||||
allocated = torch.cuda.memory_allocated()
|
||||
allocated = torch.accelerator.memory_allocated()
|
||||
assert allocated < 50 * 1024 * 1024
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from unittest.mock import mock_open, patch
|
||||
|
||||
from vllm.platforms import _is_amd_zen_cpu
|
||||
|
||||
|
||||
def test_is_amd_zen_cpu_detects_amd_with_avx512():
|
||||
cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2 avx512f avx512bw"
|
||||
with (
|
||||
patch("os.path.exists", return_value=True),
|
||||
patch("builtins.open", mock_open(read_data=cpuinfo)),
|
||||
):
|
||||
assert _is_amd_zen_cpu()
|
||||
|
||||
|
||||
def test_is_amd_zen_cpu_returns_false_for_amd_without_avx512():
|
||||
cpuinfo = "vendor_id: AuthenticAMD\nflags: avx avx2"
|
||||
with (
|
||||
patch("os.path.exists", return_value=True),
|
||||
patch("builtins.open", mock_open(read_data=cpuinfo)),
|
||||
):
|
||||
assert not _is_amd_zen_cpu()
|
||||
|
||||
|
||||
def test_is_amd_zen_cpu_returns_false_for_intel_with_avx512():
|
||||
cpuinfo = "vendor_id: GenuineIntel\nflags: avx avx2 avx512f"
|
||||
with (
|
||||
patch("os.path.exists", return_value=True),
|
||||
patch("builtins.open", mock_open(read_data=cpuinfo)),
|
||||
):
|
||||
assert not _is_amd_zen_cpu()
|
||||
|
||||
|
||||
def test_is_amd_zen_cpu_returns_false_when_cpuinfo_missing():
|
||||
with patch("os.path.exists", return_value=False):
|
||||
assert not _is_amd_zen_cpu()
|
||||
@@ -560,19 +560,23 @@ def test_streaming_empty_tool_call(glm4_moe_tool_parser, mock_request):
|
||||
assert glm4_moe_tool_parser.current_tool_id == -1
|
||||
|
||||
|
||||
def test_streaming_prev_tool_call_arr_finalization(glm4_moe_tool_parser, mock_request):
|
||||
def test_streaming_prev_tool_call_arr_updates(glm4_moe_tool_parser, mock_request):
|
||||
"""Test that prev_tool_call_arr contains parsed dict after tool call."""
|
||||
_reset_streaming_state(glm4_moe_tool_parser)
|
||||
|
||||
# Stream a complete tool call
|
||||
name_only = {"name": "get_weather", "arguments": {}}
|
||||
name_and_args = {"name": "get_weather", "arguments": {"city": "Beijing"}}
|
||||
chunks = [
|
||||
"<tool_call>get_weather\n",
|
||||
"<arg_key>city</arg_key>",
|
||||
"<arg_value>Beijing</arg_value>",
|
||||
"</tool_call>",
|
||||
# Delta, expected streamed_args_for_tool, expected prev_tool_call_arr
|
||||
("<tool_call>get_weather\n", "", name_only),
|
||||
("<arg_key>city</arg_key>", "", name_only),
|
||||
("<arg_value>Beijing</arg_value>", '{"city": "Beijing"', name_only),
|
||||
# Note: arguments are only updated when the tool call is complete.
|
||||
("</tool_call>", '{"city": "Beijing"}', name_and_args),
|
||||
]
|
||||
|
||||
for chunk in chunks:
|
||||
for chunk, exp_streamed, exp_prev_tc in chunks:
|
||||
glm4_moe_tool_parser.extract_tool_calls_streaming(
|
||||
previous_text="",
|
||||
current_text="",
|
||||
@@ -582,6 +586,8 @@ def test_streaming_prev_tool_call_arr_finalization(glm4_moe_tool_parser, mock_re
|
||||
delta_token_ids=[],
|
||||
request=mock_request,
|
||||
)
|
||||
assert glm4_moe_tool_parser.streamed_args_for_tool[0] == exp_streamed
|
||||
assert glm4_moe_tool_parser.prev_tool_call_arr[0] == exp_prev_tc
|
||||
|
||||
# After the tool call completes, prev_tool_call_arr should have parsed dict
|
||||
assert len(glm4_moe_tool_parser.prev_tool_call_arr) == 1
|
||||
@@ -592,6 +598,12 @@ def test_streaming_prev_tool_call_arr_finalization(glm4_moe_tool_parser, mock_re
|
||||
assert isinstance(args, dict), f"Expected dict, got {type(args)}"
|
||||
assert args.get("city") == "Beijing"
|
||||
|
||||
# Test equivalence of prev_tool_call_arr and streamed_args_for_tool
|
||||
# Simulates logic in chat_completion/serving.py:chat_completion_stream_generator
|
||||
tool_call_json = json.dumps(tool_entry.get("arguments", {}))
|
||||
streamed_content = glm4_moe_tool_parser.streamed_args_for_tool[0]
|
||||
assert tool_call_json.startswith(streamed_content)
|
||||
|
||||
|
||||
def test_streaming_multiple_tool_calls_sequential(glm4_moe_tool_parser, mock_request):
|
||||
"""Test streaming multiple sequential tool calls."""
|
||||
|
||||
@@ -11,6 +11,7 @@ from vllm.transformers_utils.gguf_utils import (
|
||||
split_remote_gguf,
|
||||
)
|
||||
from vllm.transformers_utils.utils import (
|
||||
is_azure,
|
||||
is_cloud_storage,
|
||||
is_gcs,
|
||||
is_s3,
|
||||
@@ -31,9 +32,17 @@ def test_is_s3():
|
||||
assert not is_s3("nfs://nfs-fqdn.local")
|
||||
|
||||
|
||||
def test_is_azure():
|
||||
assert is_azure("az://model-container/path")
|
||||
assert not is_azure("s3://model-path/path-to-model")
|
||||
assert not is_azure("/unix/local/path")
|
||||
assert not is_azure("nfs://nfs-fqdn.local")
|
||||
|
||||
|
||||
def test_is_cloud_storage():
|
||||
assert is_cloud_storage("gs://model-path")
|
||||
assert is_cloud_storage("s3://model-path/path-to-model")
|
||||
assert is_cloud_storage("az://model-container/path")
|
||||
assert not is_cloud_storage("/unix/local/path")
|
||||
assert not is_cloud_storage("nfs://nfs-fqdn.local")
|
||||
|
||||
|
||||
@@ -122,6 +122,12 @@ ROCM_EXTRA_ARGS = (
|
||||
if current_platform.is_rocm()
|
||||
else []
|
||||
)
|
||||
# Python-API equivalent of ROCM_EXTRA_ARGS for use with EngineArgs kwargs.
|
||||
ROCM_ENGINE_KWARGS: dict = (
|
||||
{"enable_prefix_caching": False, "max_num_seqs": 1}
|
||||
if current_platform.is_rocm()
|
||||
else {}
|
||||
)
|
||||
|
||||
|
||||
class RemoteVLLMServer:
|
||||
|
||||
@@ -29,7 +29,7 @@ def test_memory_profiling():
|
||||
def measure_current_non_torch():
|
||||
free, total = torch.cuda.mem_get_info()
|
||||
current_used = total - free
|
||||
current_torch = torch.cuda.memory_reserved()
|
||||
current_torch = torch.accelerator.memory_reserved()
|
||||
current_non_torch = current_used - current_torch
|
||||
return current_non_torch
|
||||
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for GDNAttentionMetadataBuilder.build() — specifically the
|
||||
reclassification of non-spec decodes as prefills when spec decodes exist.
|
||||
Covers the fix for https://github.com/vllm-project/vllm/issues/34845.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.v1.attention.utils import (
|
||||
BatchSpec,
|
||||
create_common_attn_metadata,
|
||||
create_vllm_config,
|
||||
)
|
||||
from vllm.config import SpeculativeConfig
|
||||
from vllm.v1.attention.backends.gdn_attn import (
|
||||
GDNAttentionMetadata,
|
||||
GDNAttentionMetadataBuilder,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import MambaSpec
|
||||
|
||||
BLOCK_SIZE = 16
|
||||
DEVICE = torch.device("cpu")
|
||||
|
||||
|
||||
@dataclass
|
||||
class GDNBuildTestCase:
|
||||
"""Specification for a GDN metadata builder classification test."""
|
||||
|
||||
seq_lens: list[int]
|
||||
query_lens: list[int]
|
||||
num_decode_draft_tokens: list[int] | None # None = no spec config
|
||||
num_speculative_tokens: int
|
||||
expected_num_decodes: int
|
||||
expected_num_prefills: int
|
||||
expected_num_prefill_tokens: int
|
||||
expected_num_spec_decodes: int
|
||||
|
||||
|
||||
GDN_BUILD_TEST_CASES = {
|
||||
# The original #34845 crash: non-spec query_len=1 + spec decode
|
||||
"mixed_decode_and_spec_decode": GDNBuildTestCase(
|
||||
seq_lens=[65, 20],
|
||||
query_lens=[1, 3],
|
||||
num_decode_draft_tokens=[-1, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=1,
|
||||
expected_num_prefill_tokens=1,
|
||||
expected_num_spec_decodes=1,
|
||||
),
|
||||
# All requests are spec decodes — no reclassification needed
|
||||
"pure_spec_decode": GDNBuildTestCase(
|
||||
seq_lens=[50, 30],
|
||||
query_lens=[3, 3],
|
||||
num_decode_draft_tokens=[2, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=0,
|
||||
expected_num_prefill_tokens=0,
|
||||
expected_num_spec_decodes=2,
|
||||
),
|
||||
# No speculative config at all — standard decode path
|
||||
"pure_regular_decode": GDNBuildTestCase(
|
||||
seq_lens=[40, 30, 20],
|
||||
query_lens=[1, 1, 1],
|
||||
num_decode_draft_tokens=None,
|
||||
num_speculative_tokens=0,
|
||||
expected_num_decodes=3,
|
||||
expected_num_prefills=0,
|
||||
expected_num_prefill_tokens=0,
|
||||
expected_num_spec_decodes=0,
|
||||
),
|
||||
# Multi-token prefill alongside spec decode — no decode to reclassify
|
||||
"spec_decode_with_real_prefill": GDNBuildTestCase(
|
||||
seq_lens=[100, 20],
|
||||
query_lens=[50, 3],
|
||||
num_decode_draft_tokens=[-1, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=1,
|
||||
expected_num_prefill_tokens=50,
|
||||
expected_num_spec_decodes=1,
|
||||
),
|
||||
# All three types in one batch — decode gets reclassified
|
||||
"prefill_decode_and_spec_decode": GDNBuildTestCase(
|
||||
seq_lens=[100, 65, 20],
|
||||
query_lens=[50, 1, 3],
|
||||
num_decode_draft_tokens=[-1, -1, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=2,
|
||||
expected_num_prefill_tokens=51,
|
||||
expected_num_spec_decodes=1,
|
||||
),
|
||||
# Multiple non-spec query_len=1 requests all reclassified
|
||||
"multiple_decodes_reclassified": GDNBuildTestCase(
|
||||
seq_lens=[40, 50, 60, 20],
|
||||
query_lens=[1, 1, 1, 3],
|
||||
num_decode_draft_tokens=[-1, -1, -1, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=3,
|
||||
expected_num_prefill_tokens=3,
|
||||
expected_num_spec_decodes=1,
|
||||
),
|
||||
# Zero-length padded sequence excluded from counts
|
||||
"zero_length_padding_with_spec": GDNBuildTestCase(
|
||||
seq_lens=[16, 65, 20],
|
||||
query_lens=[0, 1, 3],
|
||||
num_decode_draft_tokens=[-1, -1, 2],
|
||||
num_speculative_tokens=2,
|
||||
expected_num_decodes=0,
|
||||
expected_num_prefills=1,
|
||||
expected_num_prefill_tokens=1,
|
||||
expected_num_spec_decodes=1,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _create_gdn_builder(
|
||||
num_speculative_tokens: int = 0,
|
||||
) -> GDNAttentionMetadataBuilder:
|
||||
"""Create a GDNAttentionMetadataBuilder with minimal config."""
|
||||
vllm_config = create_vllm_config(block_size=BLOCK_SIZE)
|
||||
if num_speculative_tokens > 0:
|
||||
vllm_config.speculative_config = SpeculativeConfig(
|
||||
method="ngram",
|
||||
num_speculative_tokens=num_speculative_tokens,
|
||||
)
|
||||
mamba_spec = MambaSpec(
|
||||
block_size=BLOCK_SIZE,
|
||||
shapes=((16, 64),),
|
||||
dtypes=(torch.float16,),
|
||||
)
|
||||
return GDNAttentionMetadataBuilder(
|
||||
kv_cache_spec=mamba_spec,
|
||||
layer_names=["layer.0"],
|
||||
vllm_config=vllm_config,
|
||||
device=DEVICE,
|
||||
)
|
||||
|
||||
|
||||
def _build(
|
||||
builder: GDNAttentionMetadataBuilder,
|
||||
batch_spec: BatchSpec,
|
||||
num_decode_draft_tokens: list[int] | None = None,
|
||||
) -> GDNAttentionMetadata:
|
||||
"""Build GDN attention metadata, optionally with spec-decode kwargs."""
|
||||
common = create_common_attn_metadata(batch_spec, BLOCK_SIZE, DEVICE)
|
||||
kwargs: dict = {}
|
||||
if num_decode_draft_tokens is not None:
|
||||
kwargs["num_decode_draft_tokens_cpu"] = torch.tensor(
|
||||
num_decode_draft_tokens, dtype=torch.int32
|
||||
)
|
||||
kwargs["num_accepted_tokens"] = torch.ones(
|
||||
batch_spec.batch_size, dtype=torch.int32, device=DEVICE
|
||||
)
|
||||
return builder.build(common_prefix_len=0, common_attn_metadata=common, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"test_case", GDN_BUILD_TEST_CASES.values(), ids=GDN_BUILD_TEST_CASES.keys()
|
||||
)
|
||||
def test_gdn_build_classification(test_case: GDNBuildTestCase):
|
||||
"""Test that GDN metadata builder classifies requests correctly."""
|
||||
builder = _create_gdn_builder(test_case.num_speculative_tokens)
|
||||
batch = BatchSpec(seq_lens=test_case.seq_lens, query_lens=test_case.query_lens)
|
||||
meta = _build(builder, batch, test_case.num_decode_draft_tokens)
|
||||
|
||||
assert meta.num_decodes == test_case.expected_num_decodes
|
||||
assert meta.num_prefills == test_case.expected_num_prefills
|
||||
assert meta.num_prefill_tokens == test_case.expected_num_prefill_tokens
|
||||
assert meta.num_spec_decodes == test_case.expected_num_spec_decodes
|
||||
|
||||
|
||||
def test_has_initial_state_after_reclassification():
|
||||
"""After reclassification, num_prefills > 0 so the prefill kernel path
|
||||
should compute has_initial_state. For the reclassified request with
|
||||
context_lens > 0, the corresponding entry must be True."""
|
||||
builder = _create_gdn_builder(num_speculative_tokens=2)
|
||||
batch = BatchSpec(seq_lens=[65, 20], query_lens=[1, 3])
|
||||
meta = _build(builder, batch, num_decode_draft_tokens=[-1, 2])
|
||||
|
||||
assert meta.num_prefills > 0, "reclassification should produce prefills"
|
||||
assert meta.has_initial_state is not None
|
||||
# req0 has context_lens = 65 - 1 = 64 > 0, so has_initial_state[0] = True
|
||||
assert meta.has_initial_state[0].item() is True
|
||||
@@ -7,7 +7,10 @@ from typing import Any
|
||||
import pytest
|
||||
import torch._dynamo.config as dynamo_config
|
||||
|
||||
from tests.utils import large_gpu_mark, single_gpu_only
|
||||
from tests.utils import (
|
||||
large_gpu_mark,
|
||||
single_gpu_only,
|
||||
)
|
||||
from vllm import SamplingParams
|
||||
from vllm.logprobs import Logprob
|
||||
from vllm.platforms import current_platform
|
||||
@@ -150,6 +153,7 @@ def test_with_eagle3_spec_decoding(sample_json_schema, monkeypatch: pytest.Monke
|
||||
run_tests(monkeypatch, MTP_MODEL, test_configs, test_sampling_params)
|
||||
|
||||
|
||||
@pytest.mark.flaky(reruns=2, only_on=current_platform.is_rocm())
|
||||
def test_with_ngram_gpu_spec_decoding(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Test ngram_gpu speculative decoding with different configurations.
|
||||
|
||||
@@ -202,7 +206,6 @@ def run_tests(
|
||||
with monkeypatch.context() as m:
|
||||
# lock matmul precision to full FP32 (IEEE)
|
||||
m.setenv("VLLM_FLOAT32_MATMUL_PRECISION", "highest")
|
||||
# m.setenv("VLLM_BATCH_INVARIANT", "1")
|
||||
outputs: list[tuple[str, list, list]] = []
|
||||
for n, (
|
||||
test_preemption,
|
||||
@@ -351,6 +354,7 @@ def run_test(
|
||||
speculative_config=spec_config,
|
||||
disable_log_stats=False,
|
||||
attention_config=attention_config,
|
||||
enable_prefix_caching=False if current_platform.is_rocm() else None,
|
||||
**cache_arg,
|
||||
) as vllm_model:
|
||||
results = []
|
||||
|
||||
@@ -18,11 +18,19 @@ dp_ep_configs=(
|
||||
"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1)
|
||||
"DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)
|
||||
)
|
||||
hybrid_ssm_configs=(
|
||||
"ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"
|
||||
# TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.
|
||||
"ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"
|
||||
)
|
||||
|
||||
# Select config array based on DP_EP env var
|
||||
if [[ -n "${DP_EP:-}" ]]; then
|
||||
configs=("${dp_ep_configs[@]}")
|
||||
echo "DP_EP is set, using dp_ep_configs"
|
||||
elif [[ -n "${HYBRID_SSM:-}" ]]; then
|
||||
configs=("${hybrid_ssm_configs[@]}")
|
||||
echo "HYBRID_SSM is set, using hybrid_ssm_configs."
|
||||
else
|
||||
configs=("${tp_configs[@]}")
|
||||
fi
|
||||
|
||||
@@ -21,6 +21,11 @@
|
||||
# MODEL_NAME - target model (default: meta-llama/Llama-3.1-8B-Instruct)
|
||||
# NUM_SPEC_TOKENS - number of speculative tokens (default: 3)
|
||||
# GPU_MEMORY_UTILIZATION - (default: 0.7)
|
||||
# ATTENTION_BACKEND - attention backend to use
|
||||
# Default: TRITON_ATTN on ROCm, FLASH_ATTN on NVIDIA
|
||||
# ROCm options: TRITON_ATTN, ROCM_ATTN, ROCM_AITER_FA,
|
||||
# ROCM_AITER_UNIFIED_ATTN
|
||||
# NVIDIA options: FLASH_ATTN, FLASHINFER
|
||||
set -x
|
||||
|
||||
# ── Model & spec decode config ──────────────────────────────────────────
|
||||
@@ -51,6 +56,28 @@ GIT_ROOT=$(git rev-parse --show-toplevel)
|
||||
|
||||
SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
|
||||
|
||||
# ── Detect platform (NVIDIA vs ROCm) ────────────────────────────────────
|
||||
|
||||
if [[ "$SMI_BIN" == *"rocm"* ]]; then
|
||||
GPU_PLATFORM="rocm"
|
||||
GPU_DEVICE_VAR="HIP_VISIBLE_DEVICES"
|
||||
else
|
||||
GPU_PLATFORM="nvidia"
|
||||
GPU_DEVICE_VAR="CUDA_VISIBLE_DEVICES"
|
||||
fi
|
||||
echo "Detected GPU platform: ${GPU_PLATFORM} (using ${GPU_DEVICE_VAR})"
|
||||
|
||||
# ── Attention backend config ─────────────────────────────────────────────
|
||||
|
||||
if [[ -z "${ATTENTION_BACKEND:-}" ]]; then
|
||||
if [[ "$GPU_PLATFORM" == "rocm" ]]; then
|
||||
ATTENTION_BACKEND="TRITON_ATTN"
|
||||
else
|
||||
ATTENTION_BACKEND="FLASH_ATTN"
|
||||
fi
|
||||
fi
|
||||
echo "Using attention backend: ${ATTENTION_BACKEND}"
|
||||
|
||||
cleanup_instances() {
|
||||
echo ""
|
||||
echo "Cleaning up..."
|
||||
@@ -84,13 +111,16 @@ wait_for_server() {
|
||||
|
||||
# ── Resolve GPU list ─────────────────────────────────────────────────────
|
||||
|
||||
if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
|
||||
IFS=',' read -ra ALL_GPUS <<< "$CUDA_VISIBLE_DEVICES"
|
||||
# Accept either CUDA_VISIBLE_DEVICES or HIP_VISIBLE_DEVICES
|
||||
VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-${HIP_VISIBLE_DEVICES:-}}"
|
||||
|
||||
if [[ -n "${VISIBLE_DEVICES}" ]]; then
|
||||
IFS=',' read -ra ALL_GPUS <<< "$VISIBLE_DEVICES"
|
||||
else
|
||||
ALL_GPUS=()
|
||||
if [[ "$SMI_BIN" == *"nvidia"* ]]; then
|
||||
if [[ "$GPU_PLATFORM" == "nvidia" ]]; then
|
||||
num=$($SMI_BIN --query-gpu=name --format=csv,noheader | wc -l)
|
||||
elif [[ "$SMI_BIN" == *"rocm"* ]]; then
|
||||
elif [[ "$GPU_PLATFORM" == "rocm" ]]; then
|
||||
num=$($SMI_BIN -l | grep -c GPU)
|
||||
else
|
||||
num=1
|
||||
@@ -100,7 +130,7 @@ fi
|
||||
|
||||
TOTAL_GPUS_NEEDED=$(( (NUM_PREFILL_INSTANCES * PREFILLER_TP_SIZE) + (NUM_DECODE_INSTANCES * DECODER_TP_SIZE) ))
|
||||
if [[ ${#ALL_GPUS[@]} -lt $TOTAL_GPUS_NEEDED ]]; then
|
||||
echo "FAIL: Need $TOTAL_GPUS_NEEDED GPUs but only have ${#ALL_GPUS[@]} (CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-not set})"
|
||||
echo "FAIL: Need $TOTAL_GPUS_NEEDED GPUs but only have ${#ALL_GPUS[@]} (visible devices=${VISIBLE_DEVICES:-not set})"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
@@ -119,12 +149,14 @@ run_test_for_device() {
|
||||
echo "================================================================"
|
||||
echo "NixlConnector PD + Spec Decode Acceptance Test (kv_buffer_device=${kv_device})"
|
||||
echo "================================================================"
|
||||
echo "Model: ${MODEL_NAME}"
|
||||
echo "SD method: ${SD_METHOD}"
|
||||
echo "SD model: ${SD_MODEL}"
|
||||
echo "Spec tokens: ${NUM_SPEC_TOKENS}"
|
||||
echo "KV buffer device: ${kv_device}"
|
||||
echo "GPUs available: ${ALL_GPUS[*]}"
|
||||
echo "Model: ${MODEL_NAME}"
|
||||
echo "SD method: ${SD_METHOD}"
|
||||
echo "SD model: ${SD_MODEL}"
|
||||
echo "Spec tokens: ${NUM_SPEC_TOKENS}"
|
||||
echo "KV buffer device: ${kv_device}"
|
||||
echo "Attention backend: ${ATTENTION_BACKEND}"
|
||||
echo "GPU platform: ${GPU_PLATFORM}"
|
||||
echo "GPUs available: ${ALL_GPUS[*]}"
|
||||
echo "================================================================"
|
||||
|
||||
local PREFILL_HOSTS=()
|
||||
@@ -146,7 +178,8 @@ run_test_for_device() {
|
||||
local SIDE_CHANNEL_PORT=$((5559 + i))
|
||||
|
||||
echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
|
||||
CUDA_VISIBLE_DEVICES=$GPU_ID \
|
||||
env \
|
||||
${GPU_DEVICE_VAR}=$GPU_ID \
|
||||
VLLM_KV_CACHE_LAYOUT='HND' \
|
||||
UCX_NET_DEVICES=all \
|
||||
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
|
||||
@@ -159,7 +192,7 @@ run_test_for_device() {
|
||||
--tensor-parallel-size $PREFILLER_TP_SIZE \
|
||||
--kv-transfer-config "$kv_config" \
|
||||
--speculative-config "$PREFILL_SPEC_CONFIG" \
|
||||
--attention-backend FLASH_ATTN &
|
||||
--attention-backend $ATTENTION_BACKEND &
|
||||
|
||||
PREFILL_HOSTS+=("localhost")
|
||||
PREFILL_PORTS+=("$PORT")
|
||||
@@ -178,7 +211,8 @@ run_test_for_device() {
|
||||
local SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
|
||||
|
||||
echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
|
||||
CUDA_VISIBLE_DEVICES=$GPU_ID \
|
||||
env \
|
||||
${GPU_DEVICE_VAR}=$GPU_ID \
|
||||
VLLM_KV_CACHE_LAYOUT='HND' \
|
||||
UCX_NET_DEVICES=all \
|
||||
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
|
||||
@@ -191,7 +225,7 @@ run_test_for_device() {
|
||||
--tensor-parallel-size $DECODER_TP_SIZE \
|
||||
--kv-transfer-config "$kv_config" \
|
||||
--speculative-config "$DECODE_SPEC_CONFIG" \
|
||||
--attention-backend FLASH_ATTN &
|
||||
--attention-backend $ATTENTION_BACKEND &
|
||||
|
||||
DECODE_HOSTS+=("localhost")
|
||||
DECODE_PORTS+=("$PORT")
|
||||
@@ -218,7 +252,7 @@ run_test_for_device() {
|
||||
sleep 5
|
||||
|
||||
# Run test
|
||||
echo "Running spec decode acceptance test (kv_buffer_device=${kv_device})..."
|
||||
echo "Running spec decode acceptance test (kv_buffer_device=${kv_device}, backend=${ATTENTION_BACKEND})..."
|
||||
DECODE_PORT=${DECODE_PORTS[0]} \
|
||||
TEST_MODEL=$MODEL_NAME \
|
||||
python3 -m pytest -s -x "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py"
|
||||
@@ -234,4 +268,4 @@ for device in $KV_BUFFER_DEVICES; do
|
||||
run_test_for_device "$device"
|
||||
done
|
||||
|
||||
echo "=== All spec decode acceptance tests passed ==="
|
||||
echo "=== All spec decode acceptance tests passed (backend=${ATTENTION_BACKEND}) ==="
|
||||
|
||||
@@ -18,6 +18,7 @@ EXPECTED_VALUES = {
|
||||
"deepseek-ai/deepseek-vl2-tiny": 0.19,
|
||||
"deepseek-ai/DeepSeek-V2-Lite-Chat": 0.65,
|
||||
"google/gemma-3-4b-it": 0.74,
|
||||
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8": 0.84,
|
||||
}
|
||||
|
||||
SIMPLE_PROMPT = (
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
def test_mla_backend_rejects_cross_layer_kv_cache():
|
||||
"""MLA backends return identity permutation (layers dim first)
|
||||
to signal cross-layer KV cache is unsupported."""
|
||||
from vllm.model_executor.layers.attention.mla_attention import (
|
||||
MLACommonBackend,
|
||||
)
|
||||
|
||||
stride_order = MLACommonBackend.get_kv_cache_stride_order(
|
||||
include_num_layers_dimension=True
|
||||
)
|
||||
assert stride_order == (0, 1, 2, 3)
|
||||
assert stride_order[0] == 0 # layers dim first => no cross-layer
|
||||
assert MLACommonBackend.get_kv_cache_stride_order(
|
||||
include_num_layers_dimension=False
|
||||
) == (0, 1, 2)
|
||||
|
||||
|
||||
def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache():
|
||||
"""DeepseekV32Indexer returns identity permutation (layers dim first)
|
||||
to signal cross-layer KV cache is unsupported."""
|
||||
from vllm.v1.attention.backends.mla.indexer import (
|
||||
DeepseekV32IndexerBackend,
|
||||
)
|
||||
|
||||
stride_order = DeepseekV32IndexerBackend.get_kv_cache_stride_order(
|
||||
include_num_layers_dimension=True
|
||||
)
|
||||
assert stride_order == (0, 1, 2, 3)
|
||||
assert stride_order[0] == 0 # layers dim first => no cross-layer
|
||||
assert DeepseekV32IndexerBackend.get_kv_cache_stride_order(
|
||||
include_num_layers_dimension=False
|
||||
) == (0, 1, 2)
|
||||
@@ -84,10 +84,13 @@ def mock_parallel_groups():
|
||||
yield mock_group
|
||||
|
||||
|
||||
def _setup_kv_transfer_request(request, remote_host="127.0.0.1", fake_port=4789):
|
||||
def _setup_kv_transfer_request(
|
||||
request, remote_host="127.0.0.1", fake_port=4789, fake_transfer_id="0"
|
||||
):
|
||||
"""Setup KV transfer parameters for a request."""
|
||||
request.kv_transfer_params.update(
|
||||
{
|
||||
"transfer_id": fake_transfer_id,
|
||||
"remote_notify_port": fake_port,
|
||||
"remote_block_ids": None,
|
||||
"remote_host": remote_host,
|
||||
|
||||
@@ -53,7 +53,13 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
|
||||
from vllm.v1.attention.backends.utils import set_kv_cache_layout
|
||||
from vllm.v1.engine import EngineCoreRequest
|
||||
from vllm.v1.engine.output_processor import OutputProcessor
|
||||
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor
|
||||
from vllm.v1.kv_cache_interface import (
|
||||
AttentionSpec,
|
||||
FullAttentionSpec,
|
||||
KVCacheConfig,
|
||||
KVCacheGroupSpec,
|
||||
KVCacheTensor,
|
||||
)
|
||||
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
|
||||
from vllm.v1.request import RequestStatus
|
||||
from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
|
||||
@@ -332,8 +338,20 @@ def test_kv_transfer_handshake(dist_init):
|
||||
|
||||
# Prefill connector will register KV cache to populate proper handshake
|
||||
# metadata.
|
||||
# TODO this must match with values used in kv cache config
|
||||
kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
|
||||
kv_cache_groups = [
|
||||
KVCacheGroupSpec(
|
||||
["layer0", "layer1", "layer2"],
|
||||
FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=4,
|
||||
head_size=16,
|
||||
dtype=torch.float16,
|
||||
),
|
||||
)
|
||||
]
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=2, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups
|
||||
)
|
||||
prefill_connector = NixlConnector(
|
||||
vllm_config, KVConnectorRole.WORKER, kv_cache_config
|
||||
)
|
||||
@@ -437,7 +455,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
|
||||
self.kv_cache_layout = kv_cache_layout
|
||||
# Mock register_kv_caches attribute needed for tests that do not call it.
|
||||
self.src_xfer_handles_by_block_size = {self.block_size: 1}
|
||||
test_shape = self.attn_backend.get_kv_cache_shape(
|
||||
test_shape = self.attn_backends[0].get_kv_cache_shape(
|
||||
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
|
||||
)
|
||||
self.kv_topo = TpKVTopology(
|
||||
@@ -447,7 +465,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
|
||||
remote_block_size=self._block_size, # shared state
|
||||
is_mla=self.use_mla,
|
||||
total_num_kv_heads=self.model_config.get_total_num_kv_heads(),
|
||||
attn_backend=self.attn_backend,
|
||||
attn_backends=self.attn_backends,
|
||||
tensor_shape=test_shape,
|
||||
)
|
||||
|
||||
@@ -501,6 +519,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):
|
||||
# is started. We mock HND here.
|
||||
kv_cache_layout="HND",
|
||||
block_size=self.block_size,
|
||||
ssm_sizes=(0, 0),
|
||||
),
|
||||
remote_tp_rank=remote_tp_rank,
|
||||
remote_tp_size=remote_tp_size,
|
||||
@@ -951,6 +970,7 @@ class TestNixlHandshake:
|
||||
block_lens=worker.block_len_per_layer,
|
||||
kv_cache_layout=mismatched_layout,
|
||||
block_size=worker.block_size,
|
||||
ssm_sizes=(0, 0),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
@@ -1006,6 +1026,7 @@ class TestNixlHandshake:
|
||||
block_lens=[i * 2 for i in worker.block_len_per_layer],
|
||||
kv_cache_layout="HND",
|
||||
block_size=worker.block_size,
|
||||
ssm_sizes=(0, 0),
|
||||
)
|
||||
|
||||
# We don't check layout for homogeneous TP and MLA for now, as the
|
||||
@@ -1496,9 +1517,47 @@ def test_register_kv_caches(
|
||||
# test run if not mocking.
|
||||
mock_get_attn_backend.return_value = backend_cls
|
||||
mock_get_attn_backends.return_value = [backend_cls]
|
||||
num_layers = 32
|
||||
block_size = 16
|
||||
num_blocks = 8
|
||||
num_heads = 4
|
||||
head_size = 16
|
||||
|
||||
# TODO (NickLucche) the fact that connector depends on kv_cache_config for init
|
||||
# but cross-layer preference cant be inferred prior to creating kv_cache_config
|
||||
# is a bit awkward.
|
||||
dummy_connector = NixlConnector(
|
||||
vllm_config,
|
||||
KVConnectorRole.WORKER,
|
||||
make_kv_cache_config(block_size=block_size),
|
||||
)
|
||||
kv_cache_spec = FullAttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=num_heads,
|
||||
head_size=head_size,
|
||||
dtype=torch.float16,
|
||||
)
|
||||
if dummy_connector.prefer_cross_layer_blocks:
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[
|
||||
KVCacheTensor(
|
||||
size=kv_cache_spec.page_size_bytes * num_blocks,
|
||||
shared_by=["all-layers"],
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
],
|
||||
kv_cache_groups=[KVCacheGroupSpec(["all-layers"], kv_cache_spec)],
|
||||
)
|
||||
else:
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(["layer0", "layer1", "layer2"], kv_cache_spec)
|
||||
],
|
||||
)
|
||||
# Create connector
|
||||
kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
|
||||
connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
|
||||
connector.connector_worker = FakeNixlConnectorWorker(
|
||||
vllm_config,
|
||||
@@ -1526,35 +1585,6 @@ def test_register_kv_caches(
|
||||
or connector.prefer_cross_layer_blocks
|
||||
)
|
||||
if connector.prefer_cross_layer_blocks:
|
||||
num_layers = 32
|
||||
block_size = 16
|
||||
num_blocks = 8
|
||||
# Keep the fake worker's expected num_blocks in sync with the
|
||||
# cross-layer tensor we are about to register.
|
||||
worker_kv_cache_config = make_kv_cache_config(
|
||||
block_size=block_size, num_blocks=num_blocks
|
||||
)
|
||||
connector.connector_worker.kv_cache_config = worker_kv_cache_config
|
||||
connector.connector_worker.num_blocks = worker_kv_cache_config.num_blocks
|
||||
kv_cache_spec = AttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=4,
|
||||
head_size=64,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=num_blocks,
|
||||
kv_cache_tensors=[
|
||||
KVCacheTensor(
|
||||
size=kv_cache_spec.page_size_bytes * num_blocks,
|
||||
shared_by=["dummy-layer"],
|
||||
)
|
||||
for i in range(num_layers)
|
||||
],
|
||||
# allocate_uniform_kv_caches does not use this
|
||||
kv_cache_groups=[],
|
||||
)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
_, cross_layers_kv_cache, _ = (
|
||||
KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(
|
||||
@@ -1586,12 +1616,8 @@ def test_register_kv_caches(
|
||||
expected_blocks_count = 8
|
||||
|
||||
kv_caches = {"all-layers": cross_layers_kv_cache}
|
||||
|
||||
else:
|
||||
# Create test kv cache tensors using proper backend shape
|
||||
kv_cache_spec = cast(
|
||||
AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
|
||||
)
|
||||
kv_cache_shape = backend_cls.get_kv_cache_shape(
|
||||
num_blocks=kv_cache_config.num_blocks,
|
||||
block_size=kv_cache_spec.block_size,
|
||||
@@ -2261,7 +2287,7 @@ def test_compatibility_hash_validation(
|
||||
kv_cache_spec = cast(
|
||||
AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec
|
||||
)
|
||||
kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(
|
||||
kv_cache_shape = decode_worker.attn_backends[0].get_kv_cache_shape(
|
||||
num_blocks=kv_cache_config.num_blocks,
|
||||
block_size=kv_cache_spec.block_size,
|
||||
num_kv_heads=kv_cache_spec.num_kv_heads,
|
||||
@@ -2269,10 +2295,14 @@ def test_compatibility_hash_validation(
|
||||
)
|
||||
shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
|
||||
unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)
|
||||
# Build kv_caches from the actual layer names in kv_cache_config so that
|
||||
# _layer_specs lookups in register_kv_caches always find a matching key.
|
||||
layer_names = [
|
||||
name for group in kv_cache_config.kv_cache_groups for name in group.layer_names
|
||||
]
|
||||
kv_caches = {
|
||||
"layer0": shared_tensor,
|
||||
"layer1": unique_tensor,
|
||||
"layer2": shared_tensor,
|
||||
name: shared_tensor if i % 2 == 0 else unique_tensor
|
||||
for i, name in enumerate(layer_names)
|
||||
}
|
||||
decode_connector.register_kv_caches(kv_caches)
|
||||
|
||||
@@ -2312,6 +2342,7 @@ def test_compatibility_hash_validation(
|
||||
block_lens=[4096 * prefill_block_size], # slot_size * block_size
|
||||
kv_cache_layout="HND",
|
||||
block_size=prefill_block_size,
|
||||
ssm_sizes=(0, 0),
|
||||
)
|
||||
handshake_payload = NixlHandshakePayload(
|
||||
compatibility_hash=remote_hash,
|
||||
@@ -2391,7 +2422,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)
|
||||
remote_block_size=decode_worker._block_size, # shared state
|
||||
is_mla=decode_worker.use_mla,
|
||||
total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),
|
||||
attn_backend=backend,
|
||||
attn_backends=[backend],
|
||||
tensor_shape=test_shape,
|
||||
)
|
||||
|
||||
|
||||
@@ -74,6 +74,8 @@ def test_logical_to_kernel_block_ids_with_hma():
|
||||
# Simulate HMA scenario: logical block size = 32, kernel block size = 16
|
||||
# So each logical block maps to 2 kernel blocks eg [0]->[0,1]
|
||||
worker._physical_blocks_per_logical_kv_block = 2
|
||||
# FA + SW groups (neither is MambaSpec, so both get expanded)
|
||||
worker.kv_cache_config = make_kv_cache_config(block_size=16, hma_enabled=True)
|
||||
|
||||
# Test conversion: FA + SW group
|
||||
logical_block_ids = [[0, 1, 2], [3, 4]]
|
||||
@@ -201,3 +203,113 @@ def test_nixl_metadata_hma_block_ids_structure():
|
||||
assert len(req_meta.remote.block_ids) == 2
|
||||
assert list(req_meta.remote.block_ids[0]) == [10, 11, 12, 13, 14, 15, 16, 17]
|
||||
assert list(req_meta.remote.block_ids[1]) == [18, 19, 20, 21]
|
||||
|
||||
|
||||
@pytest.mark.cpu_test
|
||||
def test_get_block_descs_ids_hybrid_ssm():
|
||||
"""Test _get_block_descs_ids uses per-group strides for hybrid FA+SSM
|
||||
when ratio=1 (no kernel block size mismatch)."""
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
|
||||
NixlConnectorWorker,
|
||||
)
|
||||
|
||||
worker = object.__new__(NixlConnectorWorker)
|
||||
|
||||
num_blocks = 100
|
||||
engine_id = "test-engine"
|
||||
worker.num_regions = 2
|
||||
worker.dst_num_blocks = {engine_id: num_blocks}
|
||||
worker._has_mamba = True
|
||||
worker._is_mamba_group = [False, True]
|
||||
worker._physical_blocks_per_logical_kv_block = 1
|
||||
# num_descs = num_regions * num_blocks (no blocks_first doubling)
|
||||
worker.num_descs = 2 * num_blocks
|
||||
|
||||
fa_blocks = [3, 5]
|
||||
ssm_blocks = [1, 2]
|
||||
result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks))
|
||||
|
||||
# FA group: stride=num_blocks=100, offset=0
|
||||
# region0: [3, 5], region1: [103, 105]
|
||||
# SSM group: stride=logical_blocks=100 (=num_blocks/ratio=100/1),
|
||||
# offset=num_descs=200
|
||||
# region0: [201, 202], region1: [301, 302]
|
||||
expected = [3, 5, 103, 105, 201, 202, 301, 302]
|
||||
assert list(result) == expected, f"Expected {expected}, got {list(result)}"
|
||||
|
||||
|
||||
@pytest.mark.cpu_test
|
||||
def test_get_block_descs_ids_kernel_block_mismatch():
|
||||
"""Test _get_block_descs_ids uses different strides for FA (kernel blocks)
|
||||
vs SSM (logical blocks) when ratio > 1."""
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
|
||||
NixlConnectorWorker,
|
||||
)
|
||||
|
||||
worker = object.__new__(NixlConnectorWorker)
|
||||
|
||||
ratio = 4
|
||||
logical_blocks = 100
|
||||
num_blocks = logical_blocks * ratio # 400 kernel blocks
|
||||
engine_id = "test-engine"
|
||||
worker.num_regions = 2
|
||||
worker.dst_num_blocks = {engine_id: num_blocks}
|
||||
worker._has_mamba = True
|
||||
worker._is_mamba_group = [False, True]
|
||||
worker._physical_blocks_per_logical_kv_block = ratio
|
||||
worker.num_descs = 2 * num_blocks # 800
|
||||
|
||||
fa_blocks = [3, 7] # kernel-level block IDs
|
||||
ssm_blocks = [1, 2] # logical block IDs
|
||||
result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks))
|
||||
|
||||
# FA group: stride=num_blocks=400, offset=0
|
||||
# region0: [3, 7], region1: [403, 407]
|
||||
# SSM group: stride=logical_blocks=400//4=100, offset=num_descs=800
|
||||
# region0: [801, 802], region1: [901, 902]
|
||||
expected = [3, 7, 403, 407, 801, 802, 901, 902]
|
||||
assert list(result) == expected, f"Expected {expected}, got {list(result)}"
|
||||
|
||||
|
||||
@pytest.mark.cpu_test
|
||||
def test_nixl_metadata_hybrid_ssm_block_ids():
|
||||
"""Test NixlConnectorMetadata correctly stores block IDs for FA + SSM
|
||||
groups with different block counts (kernel mismatch active)."""
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
|
||||
NixlConnectorMetadata,
|
||||
)
|
||||
|
||||
metadata = NixlConnectorMetadata()
|
||||
|
||||
# FA: 8 kernel blocks (2 logical * ratio=4), SSM: 2 logical blocks
|
||||
fa_blocks = [0, 1, 2, 3, 4, 5, 6, 7]
|
||||
ssm_blocks = [0, 1]
|
||||
|
||||
metadata.add_new_req_to_recv(
|
||||
request_id="test-req-hybrid",
|
||||
local_block_ids=(fa_blocks, ssm_blocks),
|
||||
kv_transfer_params={
|
||||
"remote_block_ids": ([10, 11, 12, 13, 14, 15, 16, 17], [20, 21]),
|
||||
"remote_engine_id": "remote-engine",
|
||||
"remote_request_id": "prefill-test-req-hybrid",
|
||||
"remote_host": "localhost",
|
||||
"remote_port": 1234,
|
||||
"tp_size": 1,
|
||||
},
|
||||
)
|
||||
|
||||
assert "test-req-hybrid" in metadata.reqs_to_recv
|
||||
req_meta = metadata.reqs_to_recv["test-req-hybrid"]
|
||||
|
||||
# Verify local block IDs: different lengths per group
|
||||
assert len(req_meta.local_block_ids) == 2
|
||||
assert list(req_meta.local_block_ids[0]) == fa_blocks
|
||||
assert list(req_meta.local_block_ids[1]) == ssm_blocks
|
||||
assert len(req_meta.local_block_ids[0]) != len(req_meta.local_block_ids[1])
|
||||
|
||||
# Verify remote block IDs: same asymmetry preserved
|
||||
assert req_meta.remote is not None
|
||||
assert len(req_meta.remote.block_ids) == 2
|
||||
assert list(req_meta.remote.block_ids[0]) == [10, 11, 12, 13, 14, 15, 16, 17]
|
||||
assert list(req_meta.remote.block_ids[1]) == [20, 21]
|
||||
assert len(req_meta.remote.block_ids[0]) != len(req_meta.remote.block_ids[1])
|
||||
|
||||
@@ -252,29 +252,22 @@ def test_propose():
|
||||
]
|
||||
|
||||
# Sampled token IDs from target model
|
||||
sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)
|
||||
|
||||
# Mock scheduler output
|
||||
mock_scheduler_output = mock.MagicMock()
|
||||
sampled_token_ids = torch.tensor(
|
||||
[42, 60], dtype=torch.int32, device=device
|
||||
).unsqueeze(-1)
|
||||
|
||||
# Call propose
|
||||
with mock.patch(
|
||||
"vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
|
||||
) as mock_has_kv:
|
||||
mock_has_kv.return_value = False
|
||||
|
||||
draft_tokens, kv_connector_output = proposer.propose(
|
||||
sampled_token_ids=sampled_token_ids,
|
||||
target_hidden_states=target_hidden_states,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
scheduler_output=mock_scheduler_output,
|
||||
slot_mappings=None,
|
||||
)
|
||||
draft_tokens = proposer.propose(
|
||||
sampled_token_ids=sampled_token_ids,
|
||||
target_hidden_states=target_hidden_states,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
slot_mappings=None,
|
||||
)
|
||||
|
||||
# Verify draft tokens match sampled tokens
|
||||
# Shape should be [batch_size, 1] for num_speculative_tokens=1
|
||||
assert draft_tokens.shape == (batch_size, 1)
|
||||
assert torch.equal(draft_tokens[:, 0], sampled_token_ids)
|
||||
assert torch.equal(draft_tokens, sampled_token_ids)
|
||||
|
||||
# Verify the model was called
|
||||
model_mock.assert_called_once()
|
||||
@@ -326,21 +319,16 @@ def test_propose_different_layer_counts(num_hidden_layers):
|
||||
for _ in range(num_hidden_layers)
|
||||
]
|
||||
|
||||
sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)
|
||||
mock_scheduler_output = mock.MagicMock()
|
||||
sampled_token_ids = torch.tensor(
|
||||
[42, 60], dtype=torch.int32, device=device
|
||||
).unsqueeze(-1)
|
||||
|
||||
with mock.patch(
|
||||
"vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
|
||||
) as mock_has_kv:
|
||||
mock_has_kv.return_value = False
|
||||
|
||||
draft_tokens, _ = proposer.propose(
|
||||
sampled_token_ids=sampled_token_ids,
|
||||
target_hidden_states=target_hidden_states,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
scheduler_output=mock_scheduler_output,
|
||||
slot_mappings=None,
|
||||
)
|
||||
draft_tokens = proposer.propose(
|
||||
sampled_token_ids=sampled_token_ids,
|
||||
target_hidden_states=target_hidden_states,
|
||||
common_attn_metadata=common_attn_metadata,
|
||||
slot_mappings=None,
|
||||
)
|
||||
|
||||
assert draft_tokens.shape == (batch_size, 1)
|
||||
assert torch.equal(draft_tokens[:, 0], sampled_token_ids)
|
||||
assert torch.equal(draft_tokens, sampled_token_ids)
|
||||
|
||||
@@ -1,173 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Unit tests for GPT-OSS structural tag support in reasoning (PR #25515)."""
|
||||
|
||||
import json
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.mcp.tool_server import ToolServer
|
||||
from vllm.reasoning.gptoss_reasoning_parser import (
|
||||
GptOssReasoningParser,
|
||||
from_builtin_tool_to_tag,
|
||||
no_func_reaonsing_tag,
|
||||
tag_with_builtin_funcs,
|
||||
)
|
||||
|
||||
|
||||
class TestGptOssReasoningParser:
|
||||
"""Test cases for GptOssReasoningParser structural tag functionality."""
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tokenizer(self):
|
||||
"""Create a mock tokenizer for testing."""
|
||||
tokenizer = Mock()
|
||||
tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5])
|
||||
tokenizer.get_vocab = Mock(return_value={"<|end|>": 6})
|
||||
return tokenizer
|
||||
|
||||
@pytest.fixture
|
||||
def reasoning_parser(self, mock_tokenizer):
|
||||
"""Create a GptOssReasoningParser instance."""
|
||||
return GptOssReasoningParser(mock_tokenizer)
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tool_server_empty(self):
|
||||
"""Create a mock ToolServer with no tools."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(return_value=False)
|
||||
return tool_server
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tool_server_with_browser(self):
|
||||
"""Create a mock ToolServer with browser tool."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=lambda tool: tool == "browser")
|
||||
return tool_server
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tool_server_with_all_tools(self):
|
||||
"""Create a mock ToolServer with all builtin tools."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(
|
||||
side_effect=lambda tool: tool in ["browser", "python", "container"]
|
||||
)
|
||||
return tool_server
|
||||
|
||||
def test_prepare_structured_tag_no_tool_server(self, reasoning_parser):
|
||||
"""Test prepare_structured_tag with no tool server."""
|
||||
result = reasoning_parser.prepare_structured_tag(None, None)
|
||||
expected = json.dumps(no_func_reaonsing_tag)
|
||||
|
||||
assert result == expected
|
||||
|
||||
# Verify the structure is correct
|
||||
parsed = json.loads(result)
|
||||
assert parsed["type"] == "structural_tag"
|
||||
assert parsed["format"]["type"] == "triggered_tags"
|
||||
assert len(parsed["format"]["tags"]) == 1
|
||||
assert parsed["format"]["tags"][0]["begin"] == "<|channel|>analysis<|message|>"
|
||||
assert parsed["format"]["triggers"] == ["<|channel|>analysis"]
|
||||
|
||||
def test_prepare_structured_tag_with_all_tools(
|
||||
self, reasoning_parser, mock_tool_server_with_all_tools
|
||||
):
|
||||
"""Test prepare_structured_tag with all builtin tools."""
|
||||
result = reasoning_parser.prepare_structured_tag(
|
||||
None, mock_tool_server_with_all_tools
|
||||
)
|
||||
parsed = json.loads(result)
|
||||
|
||||
# Should have analysis tag + tags for all 3 tools (2 tags each)
|
||||
assert len(parsed["format"]["tags"]) == 7 # 1 analysis + 6 tool tags
|
||||
|
||||
# Check all tool tags are present
|
||||
tag_begins = [tag["begin"] for tag in parsed["format"]["tags"]]
|
||||
for tool in ["browser", "python", "container"]:
|
||||
assert f"<|channel|>commentary to={tool}" in tag_begins
|
||||
assert f"<|channel|>analysis to={tool}" in tag_begins
|
||||
|
||||
def test_prepare_structured_tag_with_original_tag(self, reasoning_parser):
|
||||
"""Test prepare_structured_tag when original_tag is provided."""
|
||||
original_tag = '{"custom": "tag"}'
|
||||
result = reasoning_parser.prepare_structured_tag(original_tag, None)
|
||||
|
||||
# Should return the original tag unchanged
|
||||
assert result == original_tag
|
||||
|
||||
def test_from_builtin_tool_to_tag(self):
|
||||
"""Test from_builtin_tool_to_tag function."""
|
||||
tags = from_builtin_tool_to_tag("python")
|
||||
|
||||
assert len(tags) == 2
|
||||
assert tags[0]["begin"] == "<|channel|>commentary to=python"
|
||||
assert tags[0]["content"]["type"] == "any_text"
|
||||
assert tags[0]["end"] == "<|end|>"
|
||||
|
||||
assert tags[1]["begin"] == "<|channel|>analysis to=python"
|
||||
assert tags[1]["content"]["type"] == "any_text"
|
||||
assert tags[1]["end"] == "<|end|>"
|
||||
|
||||
def test_tag_with_builtin_funcs(self):
|
||||
"""Test tag_with_builtin_funcs function."""
|
||||
builtin_tools = ["browser", "python"]
|
||||
result = tag_with_builtin_funcs(no_func_reaonsing_tag, builtin_tools)
|
||||
|
||||
assert result["type"] == "structural_tag"
|
||||
# Should have original analysis tag + 2 tags per tool
|
||||
assert len(result["format"]["tags"]) == 5 # 1 + 2*2
|
||||
|
||||
# Should have added commentary trigger
|
||||
assert "<|channel|>commentary to=" in result["format"]["triggers"]
|
||||
assert "<|channel|>analysis" in result["format"]["triggers"]
|
||||
|
||||
def test_tag_structure_invariants(self):
|
||||
"""Test that the basic tag structure follows expected format."""
|
||||
# Test the base no_func_reaonsing_tag structure
|
||||
assert no_func_reaonsing_tag["type"] == "structural_tag"
|
||||
assert no_func_reaonsing_tag["format"]["type"] == "triggered_tags"
|
||||
assert no_func_reaonsing_tag["format"]["stop_after_first"] is False
|
||||
|
||||
# Verify analysis tag structure
|
||||
analysis_tag = no_func_reaonsing_tag["format"]["tags"][0]
|
||||
assert analysis_tag["begin"] == "<|channel|>analysis<|message|>"
|
||||
assert analysis_tag["content"]["type"] == "any_text"
|
||||
assert analysis_tag["end"] == "<|end|>"
|
||||
|
||||
def test_json_serialization_valid(
|
||||
self, reasoning_parser, mock_tool_server_with_all_tools
|
||||
):
|
||||
"""Test that all generated tags produce valid JSON."""
|
||||
# Test with no tool server
|
||||
result1 = reasoning_parser.prepare_structured_tag(None, None)
|
||||
json.loads(result1) # Should not raise
|
||||
|
||||
# Test with empty tool server
|
||||
empty_server = Mock(spec=ToolServer)
|
||||
empty_server.has_tool = Mock(return_value=False)
|
||||
result2 = reasoning_parser.prepare_structured_tag(None, empty_server)
|
||||
json.loads(result2) # Should not raise
|
||||
|
||||
# Test with tools
|
||||
result3 = reasoning_parser.prepare_structured_tag(
|
||||
None, mock_tool_server_with_all_tools
|
||||
)
|
||||
json.loads(result3) # Should not raise
|
||||
|
||||
@pytest.mark.parametrize("tool_name", ["browser", "python", "container"])
|
||||
def test_single_tool_integration(self, reasoning_parser, tool_name):
|
||||
"""Test integration with individual tools."""
|
||||
tool_server = Mock(spec=ToolServer)
|
||||
tool_server.has_tool = Mock(side_effect=lambda tool: tool == tool_name)
|
||||
|
||||
result = reasoning_parser.prepare_structured_tag(None, tool_server)
|
||||
parsed = json.loads(result)
|
||||
|
||||
# Should have 1 analysis + 2 tool-specific tags
|
||||
assert len(parsed["format"]["tags"]) == 3
|
||||
|
||||
tag_begins = [tag["begin"] for tag in parsed["format"]["tags"]]
|
||||
assert f"<|channel|>commentary to={tool_name}" in tag_begins
|
||||
assert f"<|channel|>analysis to={tool_name}" in tag_begins
|
||||
@@ -8,7 +8,7 @@ import regex as re
|
||||
# Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`
|
||||
# --------------------------------------------------------------------------- #
|
||||
_TORCH_CUDA_PATTERNS = [
|
||||
r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|set_device|device\()\b",
|
||||
r"\btorch\.cuda\.(empty_cache|synchronize|device_count|current_device|memory_reserved|memory_allocated|max_memory_allocated|max_memory_reserved|reset_peak_memory_stats|memory_stats|set_device|device\()\b",
|
||||
r"\bwith\storch\.cuda\.device\b",
|
||||
]
|
||||
|
||||
|
||||
+85
-21
@@ -29,6 +29,81 @@ else:
|
||||
from torch.library import impl_abstract as register_fake
|
||||
|
||||
|
||||
# scaled_fp4_quant functional + out variant for torch.compile buffer management
|
||||
|
||||
|
||||
def create_fp4_scale_tensor(
|
||||
m: int,
|
||||
n: int,
|
||||
device: torch.device,
|
||||
is_sf_swizzled_layout: bool,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Allocate the output scale tensor for scaled_fp4_quant.
|
||||
|
||||
When is_sf_swizzled_layout=True, we use rounded values to store the
|
||||
swizzled scales. Due to the requirement of the Tensor Core, the minimum
|
||||
tile is 128x4 for the scales. So, we first pad the scales to multiples
|
||||
of 128 (rows) and 4 (cols). Then, the scales (in float8_e4m3fn) are
|
||||
packed into an int32 for every 4 values. More:
|
||||
https://docs.nvidia.com/cuda/parallel-thread-execution/
|
||||
#tcgen05-mma-scale-factor-b-layout-4x
|
||||
"""
|
||||
from vllm.utils.math_utils import round_up
|
||||
|
||||
block_size = 16
|
||||
if is_sf_swizzled_layout:
|
||||
rounded_m = round_up(m, 128)
|
||||
scale_n = n // block_size
|
||||
rounded_n = round_up(scale_n, 4)
|
||||
return torch.empty(
|
||||
(rounded_m, rounded_n // 4), device=device, dtype=torch.int32
|
||||
)
|
||||
else:
|
||||
return torch.empty((m, n // block_size), device=device, dtype=torch.uint8)
|
||||
|
||||
|
||||
def create_fp4_output_tensors(
|
||||
m: int,
|
||||
n: int,
|
||||
device: torch.device,
|
||||
is_sf_swizzled_layout: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Allocate both output tensors for scaled_fp4_quant:
|
||||
(quantized_output, output_scale).
|
||||
|
||||
Must match the C++ scaled_fp4_quant_func allocation exactly.
|
||||
"""
|
||||
output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)
|
||||
output_scale = create_fp4_scale_tensor(m, n, device, is_sf_swizzled_layout)
|
||||
return output, output_scale
|
||||
|
||||
|
||||
if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"):
|
||||
|
||||
@register_fake("_C::scaled_fp4_quant")
|
||||
def _scaled_fp4_quant_fake(
|
||||
input: torch.Tensor,
|
||||
input_scale: torch.Tensor,
|
||||
is_sf_swizzled_layout: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
n = input.shape[-1]
|
||||
m = input.numel() // n
|
||||
return create_fp4_output_tensors(m, n, input.device, is_sf_swizzled_layout)
|
||||
|
||||
@register_fake("_C::scaled_fp4_quant.out")
|
||||
def _scaled_fp4_quant_out_fake(
|
||||
input: torch.Tensor,
|
||||
input_scale: torch.Tensor,
|
||||
is_sf_swizzled_layout: bool,
|
||||
*,
|
||||
output: torch.Tensor,
|
||||
output_scale: torch.Tensor,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
|
||||
# page attention ops
|
||||
def paged_attention_v1(
|
||||
out: torch.Tensor,
|
||||
@@ -1644,7 +1719,6 @@ def scaled_fp4_quant(
|
||||
input = input.reshape(other_dims, input.shape[-1])
|
||||
m, n = input.shape
|
||||
block_size = 16
|
||||
device = input.device
|
||||
|
||||
assert n % block_size == 0, f"last dim has to be multiple of 16, but got {n}."
|
||||
assert input.dtype in (torch.float16, torch.bfloat16), (
|
||||
@@ -1658,26 +1732,16 @@ def scaled_fp4_quant(
|
||||
input, input_global_scale
|
||||
)
|
||||
else:
|
||||
# Two fp4 values will be packed into an uint8.
|
||||
output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)
|
||||
if is_sf_swizzled_layout:
|
||||
# We use the rounded values to store the swizzled values. Due to the
|
||||
# requirement of the Tensor Core, the minimum tile is 128x4 for the scales.
|
||||
# So, we first pad the scales to multiples of 128 and 4. Then, the scales
|
||||
# (in float8_e4m3fn) are packed into an int32 for every 4 values. More:
|
||||
# https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x
|
||||
round_up = lambda x, y: (x + y - 1) // y * y
|
||||
rounded_m = round_up(m, 128)
|
||||
scale_n = n // block_size
|
||||
rounded_n = round_up(scale_n, 4)
|
||||
output_scale = torch.empty(
|
||||
(rounded_m, rounded_n // 4), device=device, dtype=torch.int32
|
||||
)
|
||||
else:
|
||||
output_scale = torch.empty((m, n // 16), device=device, dtype=torch.uint8)
|
||||
|
||||
torch.ops._C.scaled_fp4_quant(
|
||||
output, input, output_scale, input_global_scale, is_sf_swizzled_layout
|
||||
# Pre-allocate and call .out variant (same behavior as old in-place API)
|
||||
output, output_scale = create_fp4_output_tensors(
|
||||
m, n, input.device, is_sf_swizzled_layout
|
||||
)
|
||||
torch.ops._C.scaled_fp4_quant.out(
|
||||
input,
|
||||
input_global_scale,
|
||||
is_sf_swizzled_layout,
|
||||
output=output,
|
||||
output_scale=output_scale,
|
||||
)
|
||||
|
||||
output_scale = output_scale.view(torch.float8_e4m3fn)
|
||||
|
||||
@@ -8,6 +8,7 @@ from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import direct_register_custom_op
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -54,6 +55,37 @@ if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"):
|
||||
return torch.empty((M, N), dtype=input.dtype, device=input.device)
|
||||
|
||||
|
||||
def _xpu_ops_deepseek_scaling_rope_impl(
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor | None,
|
||||
offsets: torch.Tensor | None,
|
||||
cos_sin_cache: torch.Tensor | None,
|
||||
rotary_dim: int,
|
||||
is_neox_style: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
assert key is not None
|
||||
return torch.ops._xpu_C.deepseek_scaling_rope(
|
||||
positions, query, key, offsets, cos_sin_cache, rotary_dim, is_neox_style
|
||||
)
|
||||
|
||||
|
||||
def _xpu_ops_deepseek_scaling_rope_fake(
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor | None,
|
||||
offsets: torch.Tensor | None,
|
||||
cos_sin_cache: torch.Tensor | None,
|
||||
rotary_dim: int,
|
||||
is_neox_style: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return query, key
|
||||
|
||||
|
||||
# Global flag to ensure ops are registered only once
|
||||
_OPS_REGISTERED = False
|
||||
|
||||
|
||||
class xpu_ops:
|
||||
@staticmethod
|
||||
def flash_attn_varlen_func(
|
||||
@@ -402,3 +434,21 @@ class xpu_ops:
|
||||
raw_topk_indices[: topk_indices.shape[0], : topk_indices.shape[1]] = (
|
||||
topk_indices
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def register_ops_once() -> None:
|
||||
global _OPS_REGISTERED
|
||||
if not _OPS_REGISTERED:
|
||||
# register all the custom ops here
|
||||
direct_register_custom_op(
|
||||
op_name="xpu_ops_deepseek_scaling_rope",
|
||||
op_func=_xpu_ops_deepseek_scaling_rope_impl,
|
||||
mutates_args=[],
|
||||
fake_impl=_xpu_ops_deepseek_scaling_rope_fake,
|
||||
dispatch_key=current_platform.dispatch_key,
|
||||
)
|
||||
|
||||
_OPS_REGISTERED = True
|
||||
|
||||
|
||||
xpu_ops.register_ops_once()
|
||||
|
||||
@@ -373,8 +373,15 @@ class InductorStandaloneAdaptor(CompilerInterface):
|
||||
break
|
||||
|
||||
if input_fake_mode is not None:
|
||||
fake_mode_ctx: Any = patch(
|
||||
"torch._inductor.standalone_compile.FakeTensorMode",
|
||||
# Use patch.object on the actual module from sys.modules
|
||||
# because in Python <=3.10 the string-based patch() resolves
|
||||
# torch._inductor.standalone_compile to the wrapper function
|
||||
# (defined in __init__.py) instead of the module.
|
||||
import sys
|
||||
|
||||
fake_mode_ctx: Any = patch.object(
|
||||
sys.modules["torch._inductor.standalone_compile"],
|
||||
"FakeTensorMode",
|
||||
lambda *a, **kw: input_fake_mode,
|
||||
)
|
||||
else:
|
||||
|
||||
@@ -16,7 +16,11 @@ from vllm.compilation.counter import compilation_counter
|
||||
from vllm.compilation.monitor import validate_cudagraph_capturing_enabled
|
||||
from vllm.config import CUDAGraphMode, VllmConfig
|
||||
from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
|
||||
from vllm.forward_context import BatchDescriptor, get_forward_context
|
||||
from vllm.forward_context import (
|
||||
BatchDescriptor,
|
||||
get_forward_context,
|
||||
is_forward_context_available,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.offloader.base import get_offloader
|
||||
from vllm.platforms import current_platform
|
||||
@@ -224,6 +228,12 @@ class CUDAGraphWrapper:
|
||||
self.concrete_cudagraph_entries.clear()
|
||||
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any | None:
|
||||
if not is_forward_context_available():
|
||||
# No forward context means we are outside the normal
|
||||
# inference path (e.g. a vision encoder forward pass).
|
||||
# Just run the underlying function without cudagraphs.
|
||||
return self.runnable(*args, **kwargs)
|
||||
|
||||
forward_context = get_forward_context()
|
||||
batch_descriptor = forward_context.batch_descriptor
|
||||
cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode
|
||||
|
||||
@@ -148,11 +148,11 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
|
||||
result_silu_mul = self.silu_and_mul_matcher(input)
|
||||
at = auto_functionalized(
|
||||
self.QUANT_OP,
|
||||
output=result,
|
||||
input=result_silu_mul,
|
||||
output_scale=output_scale,
|
||||
input_scale=scale,
|
||||
is_sf_swizzled_layout=True,
|
||||
output=result,
|
||||
output_scale=output_scale,
|
||||
)
|
||||
return at[1], at[2]
|
||||
|
||||
|
||||
@@ -47,7 +47,7 @@ if find_spec("flashinfer"):
|
||||
pass
|
||||
|
||||
if hasattr(torch.ops._C, "scaled_fp4_quant"):
|
||||
STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.default
|
||||
STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.out
|
||||
|
||||
# Max size of the input tensor per world size per device capability
|
||||
# to use flashinfer fused allreduce
|
||||
@@ -562,11 +562,11 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
|
||||
rms = self.rmsnorm_matcher(all_reduce, weight)
|
||||
quant_out_tuple = auto_functionalized(
|
||||
STATIC_FP4_QUANT_OP,
|
||||
output=quant_result,
|
||||
input=rms,
|
||||
output_scale=output_scale,
|
||||
input_scale=input_global_scale,
|
||||
is_sf_swizzled_layout=True,
|
||||
output=quant_result,
|
||||
output_scale=output_scale,
|
||||
)
|
||||
|
||||
# quant_out, allreduce_output, output_scale
|
||||
@@ -660,11 +660,11 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
|
||||
rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
|
||||
quant_out_tuple = auto_functionalized(
|
||||
STATIC_FP4_QUANT_OP,
|
||||
output=quant_result,
|
||||
input=rms,
|
||||
output_scale=output_scale,
|
||||
input_scale=input_global_scale,
|
||||
is_sf_swizzled_layout=True,
|
||||
output=quant_result,
|
||||
output_scale=output_scale,
|
||||
)
|
||||
|
||||
# quant_out, allreduce_output, output_scale
|
||||
|
||||
@@ -250,11 +250,11 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
|
||||
)
|
||||
at2 = auto_functionalized(
|
||||
self.QUANT_OP,
|
||||
output=output_quant,
|
||||
input=attn_out_view,
|
||||
output_scale=output_scale,
|
||||
input_scale=input_scale,
|
||||
is_sf_swizzled_layout=True,
|
||||
output=output_quant,
|
||||
output_scale=output_scale,
|
||||
)
|
||||
output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
|
||||
return at2[1], output_scale_view
|
||||
|
||||
@@ -38,7 +38,7 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {
|
||||
}
|
||||
|
||||
if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
|
||||
QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.default # noqa: E501
|
||||
QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out # noqa: E501
|
||||
|
||||
if current_platform.is_cuda():
|
||||
QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501
|
||||
|
||||
@@ -63,7 +63,7 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {
|
||||
kFp8DynamicTokenSym: torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501
|
||||
}
|
||||
if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
|
||||
QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.default
|
||||
QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out
|
||||
if current_platform.is_cuda():
|
||||
QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501
|
||||
QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501
|
||||
|
||||
+11
-49
@@ -10,7 +10,6 @@ from types import CodeType
|
||||
from typing import Any, ParamSpec, TypeVar
|
||||
|
||||
import torch
|
||||
import torch._C._dynamo.guards
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm.config import CompilationMode, CUDAGraphMode, get_current_vllm_config
|
||||
@@ -24,65 +23,23 @@ R = TypeVar("R")
|
||||
P = ParamSpec("P")
|
||||
|
||||
|
||||
def _noop_add_global_state_guard(
|
||||
self: torch._C._dynamo.guards.GuardManager, *args: Any, **kwargs: Any
|
||||
) -> None:
|
||||
"""No-op to skip the GLOBAL_STATE guard entirely"""
|
||||
pass
|
||||
|
||||
|
||||
def _noop_add_torch_function_mode_stack_guard(
|
||||
self: torch._C._dynamo.guards.GuardManager, *args: Any, **kwargs: Any
|
||||
) -> None:
|
||||
"""No-op to skip the TORCH_FUNCTION_MODE_STACK guard entirely"""
|
||||
pass
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _compilation_context() -> Generator[None, None, None]:
|
||||
"""Context manager for compilation settings and patches.
|
||||
"""Context manager for compilation settings.
|
||||
|
||||
This manager:
|
||||
1. Sets higher dynamo cache limits for compilation. (Needed for
|
||||
qwen2_5_vl see test_qwen2_5_vl_evs_functionality).
|
||||
Generally a recompilation can happen whenever we use a new
|
||||
backend instance in torch.compile.
|
||||
2. Patches out add_global_state_guard to skip GLOBAL_STATE guards
|
||||
3. Patches out add_torch_function_mode_stack_guard to skip
|
||||
TORCH_FUNCTION_MODE_STACK guards.
|
||||
4. Restores everything when compilation completes
|
||||
This manager sets higher dynamo cache limits for compilation.
|
||||
(Needed for qwen2_5_vl see test_qwen2_5_vl_evs_functionality).
|
||||
Generally a recompilation can happen whenever we use a new
|
||||
backend instance in torch.compile.
|
||||
"""
|
||||
# Save original values
|
||||
original_global_state_guard = (
|
||||
torch._C._dynamo.guards.GuardManager.add_global_state_guard
|
||||
)
|
||||
original_torch_function_mode_stack_guard = (
|
||||
torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard
|
||||
)
|
||||
original_cache_size = torch._dynamo.config.cache_size_limit
|
||||
original_accumulated_cache = torch._dynamo.config.accumulated_cache_size_limit
|
||||
|
||||
try:
|
||||
# Set higher cache limits for compilation
|
||||
torch._dynamo.config.cache_size_limit = 2048
|
||||
torch._dynamo.config.accumulated_cache_size_limit = 8192
|
||||
|
||||
# Patch guard manager
|
||||
torch._C._dynamo.guards.GuardManager.add_global_state_guard = (
|
||||
_noop_add_global_state_guard
|
||||
)
|
||||
torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard = (
|
||||
_noop_add_torch_function_mode_stack_guard
|
||||
)
|
||||
yield
|
||||
finally:
|
||||
# Restore original values
|
||||
torch._C._dynamo.guards.GuardManager.add_global_state_guard = (
|
||||
original_global_state_guard
|
||||
)
|
||||
torch._C._dynamo.guards.GuardManager.add_torch_function_mode_stack_guard = (
|
||||
original_torch_function_mode_stack_guard
|
||||
)
|
||||
torch._dynamo.config.cache_size_limit = original_cache_size
|
||||
torch._dynamo.config.accumulated_cache_size_limit = original_accumulated_cache
|
||||
|
||||
@@ -155,7 +112,12 @@ class TorchCompileWithNoGuardsWrapper:
|
||||
entry.guard_type == "SHAPE_ENV" for entry in x
|
||||
]
|
||||
else:
|
||||
options["guard_filter_fn"] = lambda x: [False for _ in x]
|
||||
if hasattr(torch.compiler, "skip_all_guards_unsafe"):
|
||||
# Torch 2.10+ provides skip_all_guards_unsafe
|
||||
options["guard_filter_fn"] = torch.compiler.skip_all_guards_unsafe
|
||||
else:
|
||||
# Equivalent fallback for older PyTorch: skip all guards
|
||||
options["guard_filter_fn"] = lambda x: [False for _ in x]
|
||||
|
||||
compiled_ptr: Any = self.forward
|
||||
# Validate that unbacked dynamic shapes require VLLM_USE_BYTECODE_HOOK=False
|
||||
|
||||
@@ -30,7 +30,7 @@ class AttentionConfig:
|
||||
use_cudnn_prefill: bool = False
|
||||
"""Whether to use cudnn prefill."""
|
||||
|
||||
use_trtllm_ragged_deepseek_prefill: bool = False
|
||||
use_trtllm_ragged_deepseek_prefill: bool = True
|
||||
"""Whether to use TRTLLM ragged deepseek prefill."""
|
||||
|
||||
use_trtllm_attention: bool | None = None
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user