forked from Karylab-cklius/vllm
Compare commits
46
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a8b0f65f44 | ||
|
|
8a93cda9a9 | ||
|
|
0ab46391b7 | ||
|
|
1ff9d33535 | ||
|
|
7f65f84428 | ||
|
|
ef34592a1a | ||
|
|
f69644caf8 | ||
|
|
d37e25ffbe | ||
|
|
8517cdaf90 | ||
|
|
4e498b5e5c | ||
|
|
28ee78af54 | ||
|
|
630492da30 | ||
|
|
920bf3ec84 | ||
|
|
39dff5ff39 | ||
|
|
d7af6b34d8 | ||
|
|
bbee532988 | ||
|
|
53181384e0 | ||
|
|
a0dc7a0f36 | ||
|
|
56e5810ff1 | ||
|
|
639cbfd274 | ||
|
|
a721315488 | ||
|
|
6fdb49392e | ||
|
|
cf0d279142 | ||
|
|
5497ffbf7c | ||
|
|
9af6a5ed75 | ||
|
|
7863fff6e5 | ||
|
|
0d453e2336 | ||
|
|
3f9c0c25b3 | ||
|
|
a2e776d716 | ||
|
|
4955990f1b | ||
|
|
4b64fc2cbf | ||
|
|
5f1b313900 | ||
|
|
724ed2fc35 | ||
|
|
a51376b3f0 | ||
|
|
8415bf2cdb | ||
|
|
ac062147fa | ||
|
|
617239b70c | ||
|
|
27ae676364 | ||
|
|
17ed5e61f5 | ||
|
|
5672d100ed | ||
|
|
770e9bd6b3 | ||
|
|
9efdddca28 | ||
|
|
b1b59720b2 | ||
|
|
f9f770ca0b | ||
|
|
5cba6839e6 | ||
|
|
05d610e5cd |
@@ -8,7 +8,6 @@ run_all_patterns:
|
||||
- "CMakeLists.txt"
|
||||
- "requirements/common.txt"
|
||||
- "requirements/cuda.txt"
|
||||
- "requirements/kv_connectors.txt"
|
||||
- "requirements/build/cuda.txt"
|
||||
- "requirements/test/cuda.txt"
|
||||
- "setup.py"
|
||||
|
||||
@@ -28,7 +28,6 @@ steps:
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -42,7 +41,6 @@ steps:
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -56,7 +54,6 @@ steps:
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -70,7 +67,6 @@ steps:
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -84,7 +80,6 @@ steps:
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -98,7 +93,6 @@ steps:
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -144,7 +138,6 @@ steps:
|
||||
# re-tag to default image tag and push, just in case arm64 build fails
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"'
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 13.0"
|
||||
depends_on: ~
|
||||
@@ -167,7 +160,6 @@ steps:
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"'
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
@@ -192,7 +184,6 @@ steps:
|
||||
# re-tag to default image tag and push, just in case arm64 build fails
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"'
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
@@ -214,7 +205,6 @@ steps:
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"'
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
@@ -241,7 +231,6 @@ steps:
|
||||
- "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"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"'
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
@@ -266,7 +255,6 @@ steps:
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"'
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
@@ -292,7 +280,6 @@ steps:
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"'
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
@@ -316,7 +303,6 @@ steps:
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"'
|
||||
|
||||
- block: "Build release image for x86_64 CPU"
|
||||
key: block-cpu-release-image-build
|
||||
@@ -334,7 +320,6 @@ steps:
|
||||
- "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 ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:latest"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$(buildkite-agent meta-data get release-version)"'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -354,7 +339,6 @@ steps:
|
||||
- "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 ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:latest"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:$(buildkite-agent meta-data get release-version)"'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -372,7 +356,15 @@ steps:
|
||||
- "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"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"'
|
||||
|
||||
- label: "Annotate release workflow - CUDA 13.0"
|
||||
depends_on:
|
||||
- create-multi-arch-manifest
|
||||
id: annotate-release-workflow
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-release.sh"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 12.9"
|
||||
depends_on:
|
||||
@@ -385,7 +377,6 @@ steps:
|
||||
- "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-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"'
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on:
|
||||
@@ -398,7 +389,6 @@ steps:
|
||||
- "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"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"'
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on:
|
||||
@@ -411,7 +401,6 @@ steps:
|
||||
- "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-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 12.9 Ubuntu 24.04" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"'
|
||||
|
||||
- label: "Publish nightly multi-arch image to DockerHub"
|
||||
depends_on:
|
||||
@@ -449,6 +438,59 @@ steps:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
|
||||
- block: "Publish release images to DockerHub"
|
||||
key: block-publish-release-images
|
||||
depends_on:
|
||||
- create-multi-arch-manifest
|
||||
- create-multi-arch-manifest-cuda-12-9
|
||||
- create-multi-arch-manifest-ubuntu2404
|
||||
- create-multi-arch-manifest-cuda-12-9-ubuntu2404
|
||||
- build-rocm-release-image
|
||||
- input-release-version
|
||||
# Wait for CPU builds if their block steps were unblocked, so publish
|
||||
# doesn't race the in-progress CPU build. allow_failure lets publish
|
||||
# proceed when the operator legitimately leaves the CPU block steps
|
||||
# unblocked or the CPU build fails.
|
||||
- step: build-cpu-release-image-x86
|
||||
allow_failure: true
|
||||
- step: build-cpu-release-image-arm64
|
||||
allow_failure: true
|
||||
if: build.env("NIGHTLY") != "1"
|
||||
|
||||
- label: "Publish release images to DockerHub"
|
||||
depends_on:
|
||||
- block-publish-release-images
|
||||
key: publish-release-images-dockerhub
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/publish-release-images.sh"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
password-env: DOCKERHUB_TOKEN
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
|
||||
- group: "Publish wheels"
|
||||
key: "publish-wheels"
|
||||
steps:
|
||||
- block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
|
||||
key: block-upload-release-wheels
|
||||
depends_on:
|
||||
- input-release-version
|
||||
- build-wheels
|
||||
|
||||
- label: "Upload release wheels to PyPI"
|
||||
depends_on:
|
||||
- block-upload-release-wheels
|
||||
id: upload-release-wheels
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
|
||||
|
||||
# =============================================================================
|
||||
# ROCm Release Pipeline (x86_64 only)
|
||||
# =============================================================================
|
||||
@@ -562,7 +604,7 @@ steps:
|
||||
echo ""
|
||||
echo " Build complete - Image and wheels cached"
|
||||
fi
|
||||
|
||||
|
||||
artifact_paths:
|
||||
- "artifacts/rocm-base-wheels/*.whl"
|
||||
env:
|
||||
@@ -778,7 +820,7 @@ steps:
|
||||
|
||||
# Push to ECR
|
||||
docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
|
||||
|
||||
|
||||
echo ""
|
||||
echo " Successfully built and pushed ROCm release image"
|
||||
echo " Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
|
||||
@@ -805,60 +847,3 @@ steps:
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
|
||||
# =============================================================================
|
||||
# Publish to DockerHub and PyPI (at the end so all builds complete first)
|
||||
# =============================================================================
|
||||
|
||||
- block: "Publish release images to DockerHub"
|
||||
key: block-publish-release-images
|
||||
depends_on:
|
||||
- create-multi-arch-manifest
|
||||
- create-multi-arch-manifest-cuda-12-9
|
||||
- create-multi-arch-manifest-ubuntu2404
|
||||
- create-multi-arch-manifest-cuda-12-9-ubuntu2404
|
||||
- build-rocm-release-image
|
||||
- input-release-version
|
||||
# Wait for CPU builds if their block steps were unblocked, so publish
|
||||
# doesn't race the in-progress CPU build. allow_failure lets publish
|
||||
# proceed when the operator legitimately leaves the CPU block steps
|
||||
# unblocked or the CPU build fails.
|
||||
- step: build-cpu-release-image-x86
|
||||
allow_failure: true
|
||||
- step: build-cpu-release-image-arm64
|
||||
allow_failure: true
|
||||
if: build.env("NIGHTLY") != "1"
|
||||
|
||||
- label: "Publish release images to DockerHub"
|
||||
depends_on:
|
||||
- block-publish-release-images
|
||||
key: publish-release-images-dockerhub
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/publish-release-images.sh"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
password-env: DOCKERHUB_TOKEN
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
|
||||
- group: "Publish wheels"
|
||||
key: "publish-wheels"
|
||||
steps:
|
||||
- block: "Confirm update release wheels to PyPI (experimental, use with caution)?"
|
||||
key: block-upload-release-wheels
|
||||
depends_on:
|
||||
- input-release-version
|
||||
- build-wheels
|
||||
|
||||
- label: "Upload release wheels to PyPI"
|
||||
depends_on:
|
||||
- block-upload-release-wheels
|
||||
id: upload-release-wheels
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/upload-release-wheels-pypi.sh"
|
||||
|
||||
@@ -1,9 +0,0 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Append a build artifact line to the Buildkite annotation.
|
||||
# Usage: annotate-build-artifact.sh <label> <value>
|
||||
set -e
|
||||
echo "- **${1}**: \`${2}\`" | \
|
||||
buildkite-agent annotate --append --style 'info' --context 'release-artifacts'
|
||||
Executable
+27
@@ -0,0 +1,27 @@
|
||||
#!/bin/bash
|
||||
|
||||
set -ex
|
||||
|
||||
# Get release version, default to 1.0.0.dev for nightly/per-commit builds
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version 2>/dev/null | sed 's/^v//')
|
||||
if [ -z "${RELEASE_VERSION}" ]; then
|
||||
RELEASE_VERSION="1.0.0.dev"
|
||||
fi
|
||||
|
||||
buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
|
||||
To download the wheel (by commit):
|
||||
\`\`\`
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
|
||||
(Optional) For CUDA 12.9:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl .
|
||||
|
||||
(Optional) For CPU:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
\`\`\`
|
||||
|
||||
Docker images are published automatically by the "Publish release images to DockerHub" pipeline step.
|
||||
EOF
|
||||
@@ -39,17 +39,11 @@ fi
|
||||
|
||||
set -x # avoid printing secrets above
|
||||
|
||||
# install uv if not already available
|
||||
if ! command -v uv &> /dev/null; then
|
||||
curl -LsSf https://astral.sh/uv/install.sh | UV_VERSION=0.11.14 sh
|
||||
export PATH="$HOME/.local/bin:$PATH"
|
||||
fi
|
||||
|
||||
# install twine and sdist build prerequisites using uv with Python 3.12
|
||||
uv venv --python 3.12 /tmp/vllm-release-env
|
||||
# install twine and sdist build prerequisites from pypi
|
||||
python3 -m venv /tmp/vllm-release-env
|
||||
source /tmp/vllm-release-env/bin/activate
|
||||
uv pip install twine
|
||||
uv pip install -r requirements/build/cuda.txt
|
||||
pip install twine
|
||||
pip install -r requirements/build/cuda.txt
|
||||
python3 -m twine --version
|
||||
|
||||
# copy release wheels to local directory
|
||||
|
||||
@@ -460,7 +460,7 @@ steps:
|
||||
- tests/lora
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
|
||||
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
|
||||
|
||||
#------------------------------------------------------ mi250 · model_executor -------------------------------------------------------#
|
||||
|
||||
@@ -929,6 +929,7 @@ steps:
|
||||
- tests/tokenizers_
|
||||
- tests/reasoning
|
||||
- tests/tool_parsers
|
||||
- tests/parser
|
||||
- tests/transformers_utils
|
||||
- tests/config
|
||||
commands:
|
||||
@@ -942,6 +943,7 @@ steps:
|
||||
- pytest -v -s tokenizers_
|
||||
- pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
|
||||
- pytest -v -s tool_parsers
|
||||
- pytest -v -s parser
|
||||
- pytest -v -s transformers_utils
|
||||
- pytest -v -s config
|
||||
|
||||
@@ -1320,7 +1322,6 @@ steps:
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/openai/speech_to_text/
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
@@ -1336,7 +1337,21 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Entrypoints Integration (Speech to Text) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/speech_to_text
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (LLM) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1760,7 +1775,7 @@ steps:
|
||||
- export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
- pytest -v -s -x lora/test_chatglm3_tp.py
|
||||
- pytest -v -s -x lora/test_llama_tp.py
|
||||
- pytest -v -s -x lora/test_llm_with_multi_loras.py
|
||||
- pytest -v -s -x lora/test_qwen3_with_multi_loras.py
|
||||
- pytest -v -s -x lora/test_olmoe_tp.py
|
||||
- pytest -v -s -x lora/test_gptoss_tp.py
|
||||
- pytest -v -s -x lora/test_qwen35_densemodel_lora.py
|
||||
@@ -2766,7 +2781,6 @@ steps:
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/openai/speech_to_text/
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3) # TBD
|
||||
@@ -2782,7 +2796,21 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Entrypoints Integration (Speech to Text) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/speech_to_text
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (Pooling) # TBD
|
||||
timeout_in_minutes: 180
|
||||
|
||||
@@ -20,6 +20,7 @@ steps:
|
||||
- label: Cudagraph
|
||||
key: cudagraph
|
||||
timeout_in_minutes: 20
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- tests/v1/cudagraph
|
||||
- vllm/v1/cudagraph_dispatcher.py
|
||||
|
||||
@@ -30,6 +30,7 @@ steps:
|
||||
- label: Engine (1 GPU)
|
||||
key: engine-1-gpu
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/v1/engine/
|
||||
- tests/v1/engine/
|
||||
@@ -50,6 +51,7 @@ steps:
|
||||
- label: e2e Core (1 GPU)
|
||||
key: e2e-core-1-gpu
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/v1/
|
||||
- tests/v1/e2e/general/
|
||||
|
||||
@@ -6,17 +6,19 @@ steps:
|
||||
key: entrypoints-unit-tests
|
||||
timeout_in_minutes: 10
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/entrypoints
|
||||
- tests/entrypoints/
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/tool_parsers
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling --ignore=entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (LLM)
|
||||
key: entrypoints-integration-llm
|
||||
timeout_in_minutes: 40
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/llm
|
||||
@@ -36,6 +38,7 @@ steps:
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
@@ -44,18 +47,17 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2)
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/openai/speech_to_text/
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3)
|
||||
@@ -69,12 +71,13 @@ steps:
|
||||
- tests/entrypoints/test_chat_utils
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Entrypoints Integration (API Server 2)
|
||||
key: entrypoints-integration-api-server-2
|
||||
timeout_in_minutes: 130
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/rpc
|
||||
@@ -86,10 +89,23 @@ steps:
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/rpc
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (Speech to Text)
|
||||
key: entrypoints-integration-speech_to_text
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/speech_to_text
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/speech_to_text
|
||||
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
key: entrypoints-integration-pooling
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/pooling
|
||||
@@ -101,6 +117,7 @@ steps:
|
||||
key: entrypoints-integration-responses-api
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/entrypoints/openai/responses
|
||||
@@ -115,5 +132,5 @@ steps:
|
||||
- csrc/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/model_executor/models/whisper.py
|
||||
commands: # LMEval+Transcription WER check
|
||||
commands: # LMEval
|
||||
- pytest -s entrypoints/openai/correctness/
|
||||
|
||||
@@ -17,6 +17,7 @@ steps:
|
||||
- label: Kernels Core Operation Test
|
||||
key: kernels-core-operation-test
|
||||
timeout_in_minutes: 75
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- tests/kernels/core
|
||||
@@ -41,6 +42,7 @@ steps:
|
||||
- label: Kernels Attention Test %N
|
||||
key: kernels-attention-test
|
||||
timeout_in_minutes: 35
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- csrc/attention/
|
||||
- vllm/v1/attention
|
||||
@@ -55,6 +57,7 @@ steps:
|
||||
- label: Kernels Quantization Test %N
|
||||
key: kernels-quantization-test
|
||||
timeout_in_minutes: 90
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/
|
||||
- vllm/model_executor/layers/quantization
|
||||
@@ -66,6 +69,7 @@ steps:
|
||||
- label: Kernels MoE Test %N
|
||||
key: kernels-moe-test
|
||||
timeout_in_minutes: 25
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/cutlass_w8a8/moe/
|
||||
- csrc/moe/
|
||||
@@ -82,6 +86,7 @@ steps:
|
||||
- label: Kernels Mamba Test
|
||||
key: kernels-mamba-test
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- csrc/mamba/
|
||||
- tests/kernels/mamba
|
||||
@@ -91,6 +96,7 @@ steps:
|
||||
|
||||
- label: Kernels KDA Test
|
||||
timeout_in_minutes: 20
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/fla/ops/kda.py
|
||||
- vllm/model_executor/layers/fla/ops/chunk_delta_h.py
|
||||
@@ -105,11 +111,7 @@ steps:
|
||||
device: h100
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- cmake/external_projects/deepgemm.cmake
|
||||
- tools/install_deepgemm.sh
|
||||
- tools/build_deepgemm_C.py
|
||||
- tools/setup_deepgemm_pythons.sh
|
||||
- tools/check_wheel_deepgemm.py
|
||||
- vllm/utils/deep_gemm.py
|
||||
- vllm/model_executor/layers/fused_moe
|
||||
- vllm/model_executor/layers/quantization
|
||||
@@ -119,7 +121,6 @@ steps:
|
||||
- tests/kernels/attention/test_deepgemm_attention.py
|
||||
- tests/quantization/test_cutlass_w4a16.py
|
||||
commands:
|
||||
- python3 ../tools/check_wheel_deepgemm.py
|
||||
- pytest -v -s kernels/quantization/test_block_fp8.py
|
||||
- pytest -v -s kernels/moe/test_deepgemm.py
|
||||
- pytest -v -s kernels/moe/test_batched_deepgemm.py
|
||||
@@ -247,7 +248,7 @@ steps:
|
||||
- label: Kernels FusedMoE Layer Test (2 B200s)
|
||||
key: kernels-fusedmoe-layer-test-2-b200s
|
||||
timeout_in_minutes: 90
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/cutlass_w8a8/moe/
|
||||
|
||||
@@ -5,6 +5,7 @@ steps:
|
||||
- label: LM Eval Small Models
|
||||
key: lm-eval-small-models
|
||||
timeout_in_minutes: 75
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
@@ -40,7 +41,7 @@ steps:
|
||||
- label: LM Eval Small Models (B200)
|
||||
key: lm-eval-small-models-b200
|
||||
timeout_in_minutes: 120
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
@@ -48,6 +49,18 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
|
||||
|
||||
- label: LM Eval Large Models (B200, EP)
|
||||
key: lm-eval-large-models-b200-ep
|
||||
timeout_in_minutes: 120
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell-ep.txt
|
||||
|
||||
- label: LM Eval Qwen3.5 Models (B200)
|
||||
key: lm-eval-qwen3-5-models-b200
|
||||
timeout_in_minutes: 120
|
||||
@@ -92,7 +105,7 @@ steps:
|
||||
|
||||
- label: MoE Refactor Integration Test (B200 DP - TEMPORARY)
|
||||
key: moe-refactor-integration-test-b200-dp-temporary
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
@@ -102,6 +115,7 @@ steps:
|
||||
- label: LM Eval TurboQuant KV Cache
|
||||
key: lm-eval-turboquant-kv-cache
|
||||
timeout_in_minutes: 75
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization/turboquant/
|
||||
- vllm/v1/attention/backends/turboquant_attn.py
|
||||
@@ -127,7 +141,7 @@ steps:
|
||||
- label: GPQA Eval (GPT-OSS) (B200)
|
||||
key: gpqa-eval-gpt-oss-b200
|
||||
timeout_in_minutes: 120
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
@@ -140,6 +154,7 @@ steps:
|
||||
|
||||
- label: MRCR Eval Small Models
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- tests/evals/mrcr/
|
||||
commands:
|
||||
|
||||
@@ -5,11 +5,12 @@ steps:
|
||||
- label: LoRA %N
|
||||
key: lora
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/lora
|
||||
- tests/lora
|
||||
commands:
|
||||
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_llm_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
|
||||
- pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_qwen3_with_multi_loras.py --ignore=lora/test_olmoe_tp.py --ignore=lora/test_deepseekv2_tp.py --ignore=lora/test_gptoss_tp.py --ignore=lora/test_qwen3moe_tp.py --ignore=lora/test_qwen35_densemodel_lora.py
|
||||
parallelism: 4
|
||||
|
||||
|
||||
@@ -31,7 +32,7 @@ steps:
|
||||
# requires multi-GPU testing for validation.
|
||||
- pytest -v -s -x lora/test_chatglm3_tp.py
|
||||
- pytest -v -s -x lora/test_llama_tp.py
|
||||
- pytest -v -s -x lora/test_llm_with_multi_loras.py
|
||||
- pytest -v -s -x lora/test_qwen3_with_multi_loras.py
|
||||
- pytest -v -s -x lora/test_olmoe_tp.py
|
||||
- pytest -v -s -x lora/test_gptoss_tp.py
|
||||
- pytest -v -s -x lora/test_qwen35_densemodel_lora.py
|
||||
@@ -5,6 +5,7 @@ steps:
|
||||
- label: V1 Spec Decode
|
||||
key: v1-spec-decode
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
@@ -58,6 +59,7 @@ steps:
|
||||
- label: V1 Core + KV + Metrics
|
||||
key: v1-core-kv-metrics
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
@@ -153,6 +155,7 @@ steps:
|
||||
key: examples
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/examples"
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/entrypoints
|
||||
- vllm/multimodal
|
||||
@@ -210,7 +213,9 @@ steps:
|
||||
- label: Python-only Installation
|
||||
key: python-only-installation
|
||||
depends_on: ~
|
||||
optional: true
|
||||
timeout_in_minutes: 20
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- tests/standalone_tests/python_only_compile.sh
|
||||
- setup.py
|
||||
@@ -220,6 +225,7 @@ steps:
|
||||
- label: Async Engine, Inputs, Utils, Worker
|
||||
key: async-engine-inputs-utils-worker
|
||||
timeout_in_minutes: 50
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/assets/
|
||||
- vllm/config/
|
||||
@@ -282,6 +288,7 @@ steps:
|
||||
- tests/tokenizers_
|
||||
- tests/reasoning
|
||||
- tests/tool_parsers
|
||||
- tests/parser
|
||||
- tests/transformers_utils
|
||||
- tests/config
|
||||
device: cpu-small
|
||||
@@ -296,6 +303,7 @@ steps:
|
||||
- pytest -v -s tokenizers_
|
||||
- pytest -v -s reasoning --ignore=reasoning/test_seedoss_reasoning_parser.py --ignore=reasoning/test_glm4_moe_reasoning_parser.py
|
||||
- pytest -v -s tool_parsers
|
||||
- pytest -v -s parser
|
||||
- pytest -v -s transformers_utils
|
||||
- pytest -v -s config
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ steps:
|
||||
- label: Model Executor
|
||||
key: model-executor
|
||||
timeout_in_minutes: 35
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/config/model.py
|
||||
|
||||
@@ -5,6 +5,7 @@ steps:
|
||||
- label: Model Runner V2 Core Tests
|
||||
key: model-runner-v2-core-tests
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
@@ -29,6 +30,7 @@ steps:
|
||||
key: model-runner-v2-examples
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/examples"
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/core/sched/
|
||||
@@ -102,6 +104,7 @@ steps:
|
||||
key: model-runner-v2-spec-decode
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
|
||||
@@ -6,6 +6,7 @@ steps:
|
||||
key: basic-models-tests-initialization
|
||||
timeout_in_minutes: 45
|
||||
torch_nightly: true
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/test_initialization.py
|
||||
@@ -19,6 +20,7 @@ steps:
|
||||
- label: Basic Models Tests (Extra Initialization) %N
|
||||
key: basic-models-tests-extra-initialization
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/models/
|
||||
- tests/models/test_initialization.py
|
||||
@@ -35,6 +37,7 @@ steps:
|
||||
- label: Basic Models Tests (Other)
|
||||
key: basic-models-tests-other
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/test_terratorch.py
|
||||
@@ -61,6 +64,7 @@ steps:
|
||||
working_dir: "/vllm-workspace/"
|
||||
optional: true
|
||||
soft_fail: true
|
||||
device: h200_18gb
|
||||
commands:
|
||||
- pip install --upgrade git+https://github.com/huggingface/transformers
|
||||
- pytest -v -s tests/models/test_initialization.py
|
||||
@@ -77,6 +81,7 @@ steps:
|
||||
working_dir: "/vllm-workspace/"
|
||||
optional: true
|
||||
soft_fail: true
|
||||
device: h200_18gb
|
||||
commands:
|
||||
- pip install transformers==4.57.5
|
||||
- pytest -v -s tests/models/test_initialization.py
|
||||
|
||||
@@ -5,6 +5,7 @@ steps:
|
||||
- label: Language Models Tests (Standard)
|
||||
key: language-models-tests-standard
|
||||
timeout_in_minutes: 25
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/language
|
||||
@@ -18,6 +19,7 @@ steps:
|
||||
- label: Language Models Tests (Extra Standard) %N
|
||||
key: language-models-tests-extra-standard
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/models/
|
||||
- tests/models/language/pooling/test_embedding.py
|
||||
@@ -35,6 +37,7 @@ steps:
|
||||
- label: Language Models Tests (Hybrid) %N
|
||||
key: language-models-tests-hybrid
|
||||
timeout_in_minutes: 75
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/language/generation
|
||||
@@ -60,6 +63,7 @@ steps:
|
||||
- label: Language Models Test (Extended Generation) # 80min
|
||||
key: language-models-test-extended-generation
|
||||
timeout_in_minutes: 110
|
||||
device: h200_18gb
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -85,6 +89,7 @@ steps:
|
||||
- label: Language Models Test (Extended Pooling) # 36min
|
||||
key: language-models-test-extended-pooling
|
||||
timeout_in_minutes: 50
|
||||
device: h200_18gb
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -40,6 +40,7 @@ steps:
|
||||
- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
|
||||
key: multi-modal-models-standard-3-llava-qwen2-vl
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
@@ -56,6 +57,7 @@ steps:
|
||||
- label: "Multi-Modal Models (Standard) 4: other + whisper"
|
||||
key: multi-modal-models-standard-4-other-whisper
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
@@ -94,6 +96,7 @@ steps:
|
||||
- label: Multi-Modal Accuracy Eval (Small Models) # 50min
|
||||
key: multi-modal-accuracy-eval-small-models
|
||||
timeout_in_minutes: 70
|
||||
device: h200_18gb
|
||||
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
|
||||
source_file_dependencies:
|
||||
- vllm/multimodal/
|
||||
@@ -105,6 +108,7 @@ steps:
|
||||
- label: Multi-Modal Models (Extended Generation 1)
|
||||
key: multi-modal-models-extended-generation-1
|
||||
optional: true
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal/generation
|
||||
@@ -122,6 +126,7 @@ steps:
|
||||
- label: Multi-Modal Models (Extended Generation 2)
|
||||
key: multi-modal-models-extended-generation-2
|
||||
optional: true
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal/generation
|
||||
@@ -132,6 +137,7 @@ steps:
|
||||
- label: Multi-Modal Models (Extended Generation 3)
|
||||
key: multi-modal-models-extended-generation-3
|
||||
optional: true
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal/generation
|
||||
|
||||
@@ -5,6 +5,7 @@ steps:
|
||||
- label: PyTorch Compilation Unit Tests
|
||||
key: pytorch-compilation-unit-tests
|
||||
timeout_in_minutes: 10
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
@@ -78,6 +79,7 @@ steps:
|
||||
- label: PyTorch Compilation Passes Unit Tests
|
||||
key: pytorch-compilation-passes-unit-tests
|
||||
timeout_in_minutes: 20
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
@@ -110,6 +112,7 @@ steps:
|
||||
- label: PyTorch Fullgraph Smoke Test
|
||||
key: pytorch-fullgraph-smoke-test
|
||||
timeout_in_minutes: 35
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
|
||||
@@ -5,6 +5,7 @@ steps:
|
||||
- label: Quantization
|
||||
key: quantization
|
||||
timeout_in_minutes: 90
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
@@ -42,6 +43,7 @@ steps:
|
||||
- label: Quantized Models Test
|
||||
key: quantized-models-test
|
||||
timeout_in_minutes: 60
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization
|
||||
- tests/models/quantization
|
||||
|
||||
@@ -5,6 +5,7 @@ steps:
|
||||
- label: Samplers Test
|
||||
key: samplers-test
|
||||
timeout_in_minutes: 75
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers
|
||||
- vllm/sampling_metadata.py
|
||||
|
||||
@@ -16,7 +16,7 @@ steps:
|
||||
- label: Spec Decode Eagle Nightly B200
|
||||
key: spec-decode-eagle-nightly-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
@@ -40,7 +40,7 @@ steps:
|
||||
- label: Spec Decode Speculators + MTP Nightly B200
|
||||
key: spec-decode-speculators-mtp-nightly-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
@@ -100,7 +100,7 @@ steps:
|
||||
|
||||
- label: Spec Decode MTP hybrid (B200)
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
|
||||
+2
-2
@@ -31,8 +31,8 @@
|
||||
/vllm/entrypoints/cli @hmellor @mgoin @DarkLight1337 @russellb
|
||||
/vllm/entrypoints/mcp @heheda12345
|
||||
/vllm/entrypoints/openai @aarnphm @chaunceyjiang @DarkLight1337 @russellb
|
||||
/vllm/entrypoints/openai/realtime @njhill
|
||||
/vllm/entrypoints/openai/speech_to_text @NickLucche
|
||||
/vllm/entrypoints/speech_to_text/realtime @njhill
|
||||
/vllm/entrypoints/speech_to_text @NickLucche
|
||||
/vllm/entrypoints/pooling @noooop
|
||||
/vllm/entrypoints/sagemaker @DarkLight1337
|
||||
/vllm/entrypoints/serve @njhill
|
||||
|
||||
@@ -9,9 +9,6 @@ PATH=${cuda_home}/bin:$PATH
|
||||
LD_LIBRARY_PATH=${cuda_home}/lib64:$LD_LIBRARY_PATH
|
||||
|
||||
# Install requirements
|
||||
if [ "$(echo $2 | cut -d. -f1)" = "12" ]; then
|
||||
sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt
|
||||
fi
|
||||
$python_executable -m pip install -r requirements/build/cuda.txt -r requirements/cuda.txt
|
||||
|
||||
# Limit the number of parallel jobs to avoid OOM
|
||||
|
||||
@@ -109,6 +109,24 @@ else()
|
||||
set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
|
||||
endif()
|
||||
|
||||
#
|
||||
# spinloop extension (pure CXX; must stay above the non-CUDA device branch so
|
||||
# CPU builds define the target before the early return)
|
||||
#
|
||||
set(VLLM_SPINLOOP_EXT_SRC "csrc/spinloop.cpp")
|
||||
set(SPINLOOP_COMPILE_FLAGS "")
|
||||
if(CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64")
|
||||
list(APPEND SPINLOOP_COMPILE_FLAGS "-mmwaitx")
|
||||
endif()
|
||||
define_extension_target(
|
||||
spinloop
|
||||
DESTINATION vllm
|
||||
LANGUAGE CXX
|
||||
SOURCES ${VLLM_SPINLOOP_EXT_SRC}
|
||||
COMPILE_FLAGS ${SPINLOOP_COMPILE_FLAGS}
|
||||
USE_SABI 3.11
|
||||
WITH_SOABI)
|
||||
|
||||
#
|
||||
# Forward the non-CUDA device extensions to external CMake scripts.
|
||||
#
|
||||
|
||||
@@ -53,7 +53,6 @@ backends:
|
||||
- FLASHINFER_MLA
|
||||
- FLASH_ATTN_MLA # Hopper only
|
||||
- FLASHMLA # Hopper only
|
||||
- TOKENSPEED_MLA # Blackwell + R1 dims + FP8 KV (use --kv-cache-dtype fp8)
|
||||
|
||||
device: "cuda:0"
|
||||
repeats: 100
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
# Compares all available MLA prefill backends:
|
||||
# FA backends: fa2, fa3, fa4 (FlashAttention versions)
|
||||
# Non-FA: flashinfer, cudnn, trtllm (Blackwell-only, require flashinfer)
|
||||
# CuTe DSL: tokenspeed (Blackwell + R1 dims, requires tokenspeed_mla)
|
||||
#
|
||||
# Uses cutlass_mla as the decode backend for impl construction
|
||||
# (only the prefill path is exercised).
|
||||
@@ -121,7 +120,6 @@ prefill_backends:
|
||||
- flashinfer
|
||||
- cudnn
|
||||
- trtllm
|
||||
- tokenspeed
|
||||
|
||||
device: "cuda:0"
|
||||
repeats: 20
|
||||
|
||||
@@ -179,27 +179,19 @@ def create_minimal_vllm_config(
|
||||
|
||||
if prefill_backend is not None:
|
||||
prefill_cfg = get_prefill_backend_config(prefill_backend)
|
||||
if prefill_cfg.get("mla_prefill_backend_enum") is not None:
|
||||
# Registry-based backends bypass the deprecated boolean flags.
|
||||
from vllm.v1.attention.backends.mla.prefill import MLAPrefillBackendEnum
|
||||
|
||||
vllm_config.attention_config.mla_prefill_backend = MLAPrefillBackendEnum[
|
||||
prefill_cfg["mla_prefill_backend_enum"]
|
||||
if prefill_cfg["flash_attn_version"] is not None:
|
||||
vllm_config.attention_config.flash_attn_version = prefill_cfg[
|
||||
"flash_attn_version"
|
||||
]
|
||||
else:
|
||||
if prefill_cfg["flash_attn_version"] is not None:
|
||||
vllm_config.attention_config.flash_attn_version = prefill_cfg[
|
||||
"flash_attn_version"
|
||||
]
|
||||
vllm_config.attention_config.disable_flashinfer_prefill = prefill_cfg[
|
||||
"disable_flashinfer_prefill"
|
||||
]
|
||||
vllm_config.attention_config.use_cudnn_prefill = prefill_cfg[
|
||||
"use_cudnn_prefill"
|
||||
]
|
||||
vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill = (
|
||||
prefill_cfg["use_trtllm_ragged_deepseek_prefill"]
|
||||
)
|
||||
vllm_config.attention_config.disable_flashinfer_prefill = prefill_cfg[
|
||||
"disable_flashinfer_prefill"
|
||||
]
|
||||
vllm_config.attention_config.use_cudnn_prefill = prefill_cfg[
|
||||
"use_cudnn_prefill"
|
||||
]
|
||||
vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill = prefill_cfg[
|
||||
"use_trtllm_ragged_deepseek_prefill"
|
||||
]
|
||||
|
||||
return vllm_config
|
||||
|
||||
@@ -231,17 +223,22 @@ _PREFILL_BACKEND_CONFIG: dict[str, dict] = {
|
||||
"use_trtllm_ragged_deepseek_prefill": False,
|
||||
},
|
||||
"flashinfer": {
|
||||
"mla_prefill_backend_enum": "FLASHINFER",
|
||||
"flash_attn_version": None,
|
||||
"disable_flashinfer_prefill": False,
|
||||
"use_cudnn_prefill": False,
|
||||
"use_trtllm_ragged_deepseek_prefill": False,
|
||||
},
|
||||
"cudnn": {
|
||||
# cuDNN prefill backend was removed; AttentionConfig raises on use.
|
||||
"mla_prefill_backend_enum": "FLASHINFER",
|
||||
"flash_attn_version": None,
|
||||
"disable_flashinfer_prefill": True,
|
||||
"use_cudnn_prefill": True,
|
||||
"use_trtllm_ragged_deepseek_prefill": False,
|
||||
},
|
||||
"trtllm": {
|
||||
"mla_prefill_backend_enum": "TRTLLM_RAGGED",
|
||||
},
|
||||
"tokenspeed": {
|
||||
"mla_prefill_backend_enum": "TOKENSPEED_MLA",
|
||||
"flash_attn_version": None,
|
||||
"disable_flashinfer_prefill": True,
|
||||
"use_cudnn_prefill": False,
|
||||
"use_trtllm_ragged_deepseek_prefill": True,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -628,21 +625,6 @@ def _create_backend_impl(
|
||||
# Create mock layer
|
||||
layer = MockLayer(device, impl=impl, kv_cache_spec=kv_cache_spec)
|
||||
|
||||
# Attach a prefill backend (MLAAttention does this in __init__; the metadata
|
||||
# builder reads layer.prefill_backend from static_forward_context).
|
||||
from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
|
||||
|
||||
prefill_backend_cls = get_mla_prefill_backend(vllm_config)
|
||||
layer.prefill_backend = prefill_backend_cls(
|
||||
num_heads=mla_dims["num_q_heads"],
|
||||
scale=(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"]) ** -0.5,
|
||||
kv_lora_rank=mla_dims["kv_lora_rank"],
|
||||
qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
|
||||
qk_rope_head_dim=mla_dims["qk_rope_head_dim"],
|
||||
v_head_dim=mla_dims["v_head_dim"],
|
||||
vllm_config=vllm_config,
|
||||
)
|
||||
|
||||
# Create builder instance if needed
|
||||
builder_instance = None
|
||||
if builder_class:
|
||||
@@ -979,6 +961,19 @@ def _run_mla_benchmark_batched(
|
||||
results = []
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
# Clear cached prefill backend detection functions so they re-evaluate
|
||||
# with the current VllmConfig. These are @functools.cache decorated and
|
||||
# would otherwise return stale results from a previous backend's config.
|
||||
from vllm.model_executor.layers.attention.mla_attention import (
|
||||
use_cudnn_prefill,
|
||||
use_flashinfer_prefill,
|
||||
use_trtllm_ragged_deepseek_prefill,
|
||||
)
|
||||
|
||||
use_flashinfer_prefill.cache_clear()
|
||||
use_cudnn_prefill.cache_clear()
|
||||
use_trtllm_ragged_deepseek_prefill.cache_clear()
|
||||
|
||||
# Create backend impl, layer, builder, and indexer (reused across benchmarks)
|
||||
impl, layer, builder_instance, indexer = _create_backend_impl(
|
||||
backend_cfg,
|
||||
@@ -990,35 +985,36 @@ def _run_mla_benchmark_batched(
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
)
|
||||
|
||||
# Verify the actual prefill backend matches what was requested. The
|
||||
# selector + impl construction already raise on misuse; here we just
|
||||
# check the resolved class against the requested name as a sanity guard.
|
||||
# Verify the actual prefill backend matches what was requested
|
||||
if prefill_backend is not None:
|
||||
expected_class = {
|
||||
"fa2": "FlashAttnPrefillBackend",
|
||||
"fa3": "FlashAttnPrefillBackend",
|
||||
"fa4": "FlashAttnPrefillBackend",
|
||||
"flashinfer": "FlashInferPrefillBackend",
|
||||
"trtllm": "TrtllmRaggedPrefillBackend",
|
||||
"tokenspeed": "TokenspeedMLAPrefillBackend",
|
||||
}.get(prefill_backend)
|
||||
actual_class = type(getattr(layer, "prefill_backend", None)).__name__
|
||||
if expected_class and actual_class != expected_class:
|
||||
raise RuntimeError(
|
||||
f"Prefill backend '{prefill_backend}' requested "
|
||||
f"{expected_class}, got {actual_class}. Check "
|
||||
f"attention_config plumbing or installed deps."
|
||||
)
|
||||
if prefill_backend in {"fa2", "fa3", "fa4"}:
|
||||
fa_version = int(prefill_backend[2:])
|
||||
actual_fa_version = getattr(
|
||||
layer.prefill_backend, "vllm_flash_attn_version", None
|
||||
)
|
||||
prefill_cfg = get_prefill_backend_config(prefill_backend)
|
||||
fa_version = prefill_cfg["flash_attn_version"]
|
||||
|
||||
if fa_version is not None:
|
||||
# FA backend: verify the impl's FA version
|
||||
actual_fa_version = getattr(impl, "vllm_flash_attn_version", None)
|
||||
if actual_fa_version != fa_version:
|
||||
raise RuntimeError(
|
||||
f"Prefill backend '{prefill_backend}' requested FA "
|
||||
f"version {fa_version}, got "
|
||||
f"{actual_fa_version} on {actual_class}."
|
||||
f"version {fa_version}, but the impl is using FA "
|
||||
f"version {actual_fa_version}. Check "
|
||||
f"vllm/v1/attention/backends/fa_utils.py."
|
||||
)
|
||||
else:
|
||||
# Non-FA backend: verify the builder picked the right path
|
||||
expected_flags = {
|
||||
"flashinfer": "_use_fi_prefill",
|
||||
"cudnn": "_use_cudnn_prefill",
|
||||
"trtllm": "_use_trtllm_ragged_prefill",
|
||||
}
|
||||
flag_name = expected_flags.get(prefill_backend)
|
||||
if flag_name and not getattr(builder_instance, flag_name, False):
|
||||
raise RuntimeError(
|
||||
f"Prefill backend '{prefill_backend}' was requested "
|
||||
f"but the metadata builder did not enable it. This "
|
||||
f"usually means a dependency is missing (e.g., "
|
||||
f"flashinfer not installed) or the platform doesn't "
|
||||
f"support it."
|
||||
)
|
||||
|
||||
# Run each benchmark with the shared impl
|
||||
|
||||
@@ -53,67 +53,48 @@ cuda_archs_loose_intersection(DEEPGEMM_ARCHS
|
||||
if(DEEPGEMM_ARCHS)
|
||||
message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}")
|
||||
|
||||
# Build _C once per interpreter in DEEPGEMM_PYTHON_INTERPRETERS (":"-
|
||||
# separated paths) so the wheel imports cleanly on every supported Python.
|
||||
# Unset → fall back to the build interpreter (editable / source builds).
|
||||
# The compile is delegated to tools/build_deepgemm_C.py and always runs
|
||||
# against the build interpreter's torch — target Pythons don't need torch.
|
||||
# Note: empty-but-set env vars are still DEFINED in cmake; treat empty as
|
||||
# unset so an empty interpreter list falls back to the build interpreter
|
||||
# rather than silently skipping the per-Python build.
|
||||
if(NOT "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}" STREQUAL "")
|
||||
string(REPLACE ":" ";" _dg_pythons "$ENV{DEEPGEMM_PYTHON_INTERPRETERS}")
|
||||
else()
|
||||
set(_dg_pythons "${Python_EXECUTABLE}")
|
||||
endif()
|
||||
message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}")
|
||||
find_package(CUDAToolkit REQUIRED)
|
||||
|
||||
# Header set fed to add_custom_command's DEPENDS so a header-only edit
|
||||
# (in upstream DeepGEMM or its vendored cutlass/fmt) re-triggers the
|
||||
# rebuild. add_custom_command does no implicit header scanning, unlike
|
||||
# add_library.
|
||||
file(GLOB_RECURSE _dg_headers
|
||||
"${deepgemm_SOURCE_DIR}/csrc/*.h"
|
||||
"${deepgemm_SOURCE_DIR}/csrc/*.hpp"
|
||||
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h"
|
||||
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp"
|
||||
"${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh")
|
||||
#
|
||||
# Build the _C pybind11 extension from DeepGEMM's C++ source.
|
||||
# This is a CXX-only module — CUDA kernels are JIT-compiled at runtime.
|
||||
#
|
||||
Python_add_library(_deep_gemm_C MODULE WITH_SOABI
|
||||
"${deepgemm_SOURCE_DIR}/csrc/python_api.cpp")
|
||||
|
||||
set(_dg_markers)
|
||||
set(_dg_seen_soabis)
|
||||
foreach(_pybin IN LISTS _dg_pythons)
|
||||
execute_process(
|
||||
COMMAND "${_pybin}" -c
|
||||
"import sysconfig; print(sysconfig.get_config_var('SOABI'))"
|
||||
OUTPUT_VARIABLE _dg_soabi
|
||||
OUTPUT_STRIP_TRAILING_WHITESPACE
|
||||
COMMAND_ERROR_IS_FATAL ANY)
|
||||
# Dedup so duplicate paths (or two paths resolving to the same CPython)
|
||||
# don't register conflicting build rules.
|
||||
if(_dg_soabi IN_LIST _dg_seen_soabis)
|
||||
continue()
|
||||
endif()
|
||||
list(APPEND _dg_seen_soabis "${_dg_soabi}")
|
||||
set(_dg_dir "${CMAKE_CURRENT_BINARY_DIR}/deepgemm_C_${_dg_soabi}")
|
||||
set(_dg_marker "${_dg_dir}/.built")
|
||||
add_custom_command(
|
||||
OUTPUT "${_dg_marker}"
|
||||
COMMAND "${Python_EXECUTABLE}"
|
||||
"${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
|
||||
"${deepgemm_SOURCE_DIR}" "${_dg_dir}" "${_pybin}"
|
||||
COMMAND "${CMAKE_COMMAND}" -E touch "${_dg_marker}"
|
||||
DEPENDS "${CMAKE_SOURCE_DIR}/tools/build_deepgemm_C.py"
|
||||
"${deepgemm_SOURCE_DIR}/csrc/python_api.cpp"
|
||||
${_dg_headers}
|
||||
COMMENT "Building DeepGEMM _C for ${_pybin}"
|
||||
VERBATIM)
|
||||
list(APPEND _dg_markers "${_dg_marker}")
|
||||
install(DIRECTORY "${_dg_dir}/"
|
||||
DESTINATION vllm/third_party/deep_gemm
|
||||
COMPONENT _deep_gemm_C
|
||||
FILES_MATCHING PATTERN "_C.cpython-*.so")
|
||||
endforeach()
|
||||
add_custom_target(_deep_gemm_C ALL DEPENDS ${_dg_markers})
|
||||
# The pybind11 module name must be _C to match DeepGEMM's Python imports.
|
||||
set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C")
|
||||
|
||||
target_compile_definitions(_deep_gemm_C PRIVATE
|
||||
"-DTORCH_EXTENSION_NAME=_C")
|
||||
|
||||
target_include_directories(_deep_gemm_C PRIVATE
|
||||
"${deepgemm_SOURCE_DIR}/csrc"
|
||||
"${deepgemm_SOURCE_DIR}/deep_gemm/include"
|
||||
"${deepgemm_SOURCE_DIR}/third-party/cutlass/include"
|
||||
"${deepgemm_SOURCE_DIR}/third-party/cutlass/tools/util/include"
|
||||
"${deepgemm_SOURCE_DIR}/third-party/fmt/include")
|
||||
|
||||
target_compile_options(_deep_gemm_C PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-O3>
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-Wno-psabi>
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-Wno-deprecated-declarations>)
|
||||
|
||||
# torch_python is required because DeepGEMM uses pybind11 type casters
|
||||
# for at::Tensor (via PYBIND11_MODULE), unlike vLLM's own extensions which
|
||||
# use torch::Library custom ops.
|
||||
find_library(TORCH_PYTHON_LIBRARY torch_python
|
||||
PATHS "${TORCH_INSTALL_PREFIX}/lib"
|
||||
REQUIRED)
|
||||
|
||||
target_link_libraries(_deep_gemm_C PRIVATE
|
||||
torch ${TORCH_LIBRARIES} "${TORCH_PYTHON_LIBRARY}"
|
||||
CUDA::cudart CUDA::nvrtc)
|
||||
|
||||
# Install the shared library into the vendored package directory
|
||||
install(TARGETS _deep_gemm_C
|
||||
LIBRARY DESTINATION vllm/third_party/deep_gemm
|
||||
COMPONENT _deep_gemm_C)
|
||||
|
||||
#
|
||||
# Vendor DeepGEMM Python package files
|
||||
|
||||
+75
-60
@@ -21,28 +21,33 @@ namespace vllm {
|
||||
|
||||
// NOTE Be EXTRA careful with raw_kv_scalar_t, for __half and __nv_bfloat16 it's
|
||||
// using u16 as the backing type.
|
||||
template <typename qk_t, bool IS_NEOX, typename raw_kv_scalar_t,
|
||||
typename cache_t, Fp8KVCacheDataType kv_dt>
|
||||
template <typename qk_t, typename cos_sin_t, bool IS_NEOX,
|
||||
typename raw_kv_scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
|
||||
__global__ void concat_and_cache_mla_rope_fused_kernel(
|
||||
const int64_t* __restrict__ positions, // [num_tokens]
|
||||
qk_t* __restrict__ q_pe, // [num_tokens, num_q_heads, rot_dim]
|
||||
qk_t* __restrict__ k_pe, // [num_tokens, rot_dim]
|
||||
const qk_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank]
|
||||
const qk_t* __restrict__ rope_cos_sin_cache, // [max_position, 2,
|
||||
// rot_dim // 2]
|
||||
const cos_sin_t* __restrict__ rope_cos_sin_cache, // [max_position, 2,
|
||||
// rot_dim // 2]
|
||||
const int rot_dim, const int64_t q_pe_stride_token,
|
||||
const int64_t q_pe_stride_head, const int64_t k_pe_stride,
|
||||
const int64_t kv_c_stride, const int num_q_heads,
|
||||
cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank +
|
||||
// rot_dim)]
|
||||
const int64_t* __restrict__ kv_cache_slot_mapping, // [num_tokens]
|
||||
const int64_t* __restrict__ slot_mapping, // [num_tokens]
|
||||
const int block_stride, const int entry_stride, const int kv_lora_rank,
|
||||
const int block_size, const float* kv_cache_quant_scale) {
|
||||
// Each thread block is responsible for one token.
|
||||
const int64_t token_idx = blockIdx.x;
|
||||
const int64_t slot_idx = slot_mapping[token_idx];
|
||||
// NOTE: slot_idx can be -1 if the token is padded
|
||||
if (slot_idx < 0) {
|
||||
return;
|
||||
}
|
||||
const int64_t pos = positions[token_idx];
|
||||
|
||||
const qk_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim;
|
||||
const cos_sin_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim;
|
||||
|
||||
const int embed_dim = rot_dim / 2;
|
||||
|
||||
@@ -54,8 +59,8 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
|
||||
|
||||
// NOTE: Would be nice to have interleaved sin/cos so we could just load
|
||||
// both at the same time.
|
||||
qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
|
||||
qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);
|
||||
qk_t cos = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx));
|
||||
qk_t sin = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim));
|
||||
|
||||
qk_t* q_pe_head_ptr =
|
||||
q_pe + token_idx * q_pe_stride_token + head_idx * q_pe_stride_head;
|
||||
@@ -81,21 +86,15 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
|
||||
q_pe_head_ptr[pair_idx_y] = y_dst;
|
||||
}
|
||||
|
||||
const int64_t slot_idx = kv_cache_slot_mapping[token_idx];
|
||||
const int64_t block_idx = slot_idx / block_size;
|
||||
const int64_t entry_idx = slot_idx % block_size;
|
||||
|
||||
// NOTE: slot_idx can be -1 if the token is padded
|
||||
if (slot_idx < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// K with 1 HEAD
|
||||
for (int i = threadIdx.x; i < embed_dim; i += blockDim.x) {
|
||||
int pair_idx = i;
|
||||
|
||||
qk_t cos = VLLM_LDG(cos_sin_ptr + pair_idx);
|
||||
qk_t sin = VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim);
|
||||
qk_t cos = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx));
|
||||
qk_t sin = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim));
|
||||
|
||||
qk_t* k_pe_head_ptr = k_pe + token_idx * k_pe_stride;
|
||||
|
||||
@@ -165,36 +164,43 @@ __global__ void concat_and_cache_mla_rope_fused_kernel(
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
|
||||
do { \
|
||||
VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \
|
||||
using qk_t = scalar_t; \
|
||||
if (rope_is_neox) { \
|
||||
vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, true, RAW_KV_T, \
|
||||
CACHE_T, KV_DTYPE> \
|
||||
<<<grid, block, 0, stream>>>( \
|
||||
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
|
||||
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
|
||||
rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
|
||||
q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
|
||||
num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
|
||||
kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
|
||||
entry_stride, kv_lora_rank, block_size, \
|
||||
kv_cache_quant_scale.data_ptr<float>()); \
|
||||
} else { \
|
||||
vllm::concat_and_cache_mla_rope_fused_kernel<qk_t, false, RAW_KV_T, \
|
||||
CACHE_T, KV_DTYPE> \
|
||||
<<<grid, block, 0, stream>>>( \
|
||||
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
|
||||
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
|
||||
rope_cos_sin_cache.data_ptr<qk_t>(), rot_dim, \
|
||||
q_pe_stride_token, q_pe_stride_head, k_pe_stride, kv_c_stride, \
|
||||
num_q_heads, reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
|
||||
kv_cache_slot_mapping.data_ptr<int64_t>(), block_stride, \
|
||||
entry_stride, kv_lora_rank, block_size, \
|
||||
kv_cache_quant_scale.data_ptr<float>()); \
|
||||
} \
|
||||
}); \
|
||||
#define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \
|
||||
do { \
|
||||
VLLM_DISPATCH_FLOATING_TYPES(q_pe.scalar_type(), "qk_scalar_type", [&] { \
|
||||
using qk_t = scalar_t; \
|
||||
VLLM_DISPATCH_FLOATING_TYPES( \
|
||||
rope_cos_sin_cache.scalar_type(), "rope_cos_sin_cache_scalar_type", \
|
||||
[&] { \
|
||||
using cos_sin_t = scalar_t; \
|
||||
if (rope_is_neox) { \
|
||||
vllm::concat_and_cache_mla_rope_fused_kernel< \
|
||||
qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \
|
||||
<<<grid, block, 0, stream>>>( \
|
||||
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
|
||||
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
|
||||
rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
|
||||
q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
|
||||
kv_c_stride, num_q_heads, \
|
||||
reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
|
||||
slot_mapping.data_ptr<int64_t>(), block_stride, \
|
||||
entry_stride, kv_lora_rank, block_size, \
|
||||
kv_cache_quant_scale.data_ptr<float>()); \
|
||||
} else { \
|
||||
vllm::concat_and_cache_mla_rope_fused_kernel< \
|
||||
qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \
|
||||
<<<grid, block, 0, stream>>>( \
|
||||
positions.data_ptr<int64_t>(), q_pe.data_ptr<qk_t>(), \
|
||||
k_pe.data_ptr<qk_t>(), kv_c.data_ptr<qk_t>(), \
|
||||
rope_cos_sin_cache.data_ptr<cos_sin_t>(), rot_dim, \
|
||||
q_pe_stride_token, q_pe_stride_head, k_pe_stride, \
|
||||
kv_c_stride, num_q_heads, \
|
||||
reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()), \
|
||||
slot_mapping.data_ptr<int64_t>(), block_stride, \
|
||||
entry_stride, kv_lora_rank, block_size, \
|
||||
kv_cache_quant_scale.data_ptr<float>()); \
|
||||
} \
|
||||
}); \
|
||||
}); \
|
||||
} while (false)
|
||||
|
||||
// Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache.
|
||||
@@ -208,43 +214,52 @@ void concat_and_cache_mla_rope_fused(
|
||||
torch::Tensor& kv_c, // [num_tokens, kv_lora_rank]
|
||||
torch::Tensor& rope_cos_sin_cache, // [max_position, rot_dim]
|
||||
bool rope_is_neox,
|
||||
torch::Tensor&
|
||||
kv_cache_slot_mapping, // [num_tokens] or [num_actual_tokens]
|
||||
torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
|
||||
torch::Tensor&
|
||||
kv_cache, // [num_blocks, block_size, (kv_lora_rank + rot_dim)]
|
||||
const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale) {
|
||||
const int64_t num_tokens = q_pe.size(0);
|
||||
// NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from
|
||||
// slot_mapping.size(0) because of padding for CUDA graphs.
|
||||
// In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
|
||||
// both include padding.
|
||||
// In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
|
||||
// since key includes padding for CUDA graphs, while slot_mapping does not.
|
||||
// In this case, slot_mapping.size(0) represents the actual number of tokens
|
||||
// before padding.
|
||||
// For compatibility with both cases, we use slot_mapping.size(0) as the
|
||||
// number of tokens.
|
||||
int num_tokens = slot_mapping.size(0);
|
||||
int num_padded_tokens = q_pe.size(0);
|
||||
TORCH_CHECK_GE(num_padded_tokens, num_tokens);
|
||||
|
||||
const int num_q_heads = q_pe.size(1);
|
||||
const int rot_dim = q_pe.size(2);
|
||||
const int kv_lora_rank = kv_c.size(1);
|
||||
|
||||
TORCH_CHECK(positions.size(0) >=
|
||||
num_tokens); // CUDA Graphs might pad this for us
|
||||
TORCH_CHECK_EQ(positions.size(0), num_padded_tokens);
|
||||
TORCH_CHECK_EQ(positions.dim(), 1);
|
||||
TORCH_CHECK_EQ(positions.scalar_type(), c10::ScalarType::Long);
|
||||
|
||||
TORCH_CHECK_EQ(q_pe.size(0), num_tokens);
|
||||
TORCH_CHECK_EQ(q_pe.dim(), 3);
|
||||
TORCH_CHECK_EQ(q_pe.size(0), num_padded_tokens);
|
||||
TORCH_CHECK_EQ(q_pe.size(1), num_q_heads);
|
||||
TORCH_CHECK_EQ(q_pe.size(2), rot_dim);
|
||||
TORCH_CHECK_EQ(q_pe.dim(), 3);
|
||||
|
||||
TORCH_CHECK_EQ(k_pe.size(0), num_tokens);
|
||||
TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
|
||||
TORCH_CHECK_EQ(k_pe.dim(), 2);
|
||||
TORCH_CHECK_EQ(k_pe.size(0), num_padded_tokens);
|
||||
TORCH_CHECK_EQ(k_pe.size(1), rot_dim);
|
||||
TORCH_CHECK_EQ(k_pe.scalar_type(), q_pe.scalar_type());
|
||||
|
||||
TORCH_CHECK_EQ(kv_c.size(0), num_tokens);
|
||||
TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
|
||||
TORCH_CHECK_EQ(kv_c.dim(), 2);
|
||||
TORCH_CHECK_EQ(kv_c.size(0), num_padded_tokens);
|
||||
TORCH_CHECK_EQ(kv_c.size(1), kv_lora_rank);
|
||||
TORCH_CHECK_EQ(kv_c.scalar_type(), q_pe.scalar_type());
|
||||
TORCH_CHECK_EQ(kv_c.dtype(), q_pe.dtype());
|
||||
|
||||
TORCH_CHECK_EQ(rope_cos_sin_cache.size(1), rot_dim);
|
||||
TORCH_CHECK_EQ(rope_cos_sin_cache.scalar_type(), q_pe.scalar_type());
|
||||
|
||||
TORCH_CHECK_EQ(kv_cache_slot_mapping.size(0), num_tokens);
|
||||
TORCH_CHECK_EQ(kv_cache_slot_mapping.scalar_type(), c10::ScalarType::Long);
|
||||
TORCH_CHECK_EQ(slot_mapping.size(0), num_tokens);
|
||||
TORCH_CHECK_EQ(slot_mapping.scalar_type(), c10::ScalarType::Long);
|
||||
|
||||
TORCH_CHECK_EQ(kv_cache.size(2), kv_lora_rank + rot_dim);
|
||||
TORCH_CHECK_EQ(kv_cache.dim(), 3);
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
#pragma once
|
||||
#include <cstdlib>
|
||||
#include <string>
|
||||
#include <cctype>
|
||||
|
||||
namespace vllm {
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
#include "scaled_mm_kernels.hpp"
|
||||
#include "scaled_mm_sm100_fp8_dispatch.cuh"
|
||||
#include "core/batch_invariant.hpp"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -13,9 +14,17 @@ void cutlass_scaled_mm_sm100_fp8(
|
||||
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
|
||||
"currently bias dtype must match output dtype ",
|
||||
out.scalar_type());
|
||||
if (vllm_is_batch_invariant()) {
|
||||
return cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue<true>(
|
||||
out, a, b, a_scales, b_scales, *bias);
|
||||
}
|
||||
return cutlass_scaled_mm_sm100_fp8_epilogue<true>(out, a, b, a_scales,
|
||||
b_scales, *bias);
|
||||
} else {
|
||||
if (vllm_is_batch_invariant()) {
|
||||
return cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue<false>(
|
||||
out, a, b, a_scales, b_scales);
|
||||
}
|
||||
return cutlass_scaled_mm_sm100_fp8_epilogue<false>(out, a, b, a_scales,
|
||||
b_scales);
|
||||
}
|
||||
|
||||
@@ -294,6 +294,34 @@ inline void cutlass_gemm_sm100_fp8_dispatch(
|
||||
}
|
||||
}
|
||||
|
||||
template <typename InType, typename OutType, bool EnableBias,
|
||||
typename... EpilogueArgs>
|
||||
inline void cutlass_gemm_sm100_fp8_batch_invariant_dispatch(
|
||||
torch::stable::Tensor& out, torch::stable::Tensor const& a,
|
||||
torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
|
||||
torch::stable::Tensor const& b_scales, EpilogueArgs&&... args) {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
STD_TORCH_CHECK(a.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
STD_TORCH_CHECK(b.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
|
||||
using Cutlass3xGemmM64SwapAB =
|
||||
typename sm100_fp8_config_M64_swap_ab<InType, OutType,
|
||||
EnableBias>::Cutlass3xGemm;
|
||||
using Cutlass3xGemmM64 =
|
||||
typename sm100_fp8_config_M64<InType, OutType, EnableBias>::Cutlass3xGemm;
|
||||
|
||||
// keep the CUTLASS config independent of M for batch invariance
|
||||
uint32_t const k = a.size(1);
|
||||
if (k < 4096) {
|
||||
return cutlass_gemm_caller_sm100_fp8<Cutlass3xGemmM64>(
|
||||
out, a, b, a_scales, b_scales, std::forward<EpilogueArgs>(args)...);
|
||||
}
|
||||
return cutlass_gemm_caller_sm100_fp8<Cutlass3xGemmM64SwapAB>(
|
||||
out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
|
||||
}
|
||||
|
||||
template <bool EnableBias, typename... EpilogueArgs>
|
||||
void cutlass_scaled_mm_sm100_fp8_epilogue(torch::stable::Tensor& out,
|
||||
torch::stable::Tensor const& a,
|
||||
@@ -320,4 +348,28 @@ void cutlass_scaled_mm_sm100_fp8_epilogue(torch::stable::Tensor& out,
|
||||
}
|
||||
}
|
||||
|
||||
template <bool EnableBias, typename... EpilogueArgs>
|
||||
void cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue(
|
||||
torch::stable::Tensor& out, torch::stable::Tensor const& a,
|
||||
torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
|
||||
torch::stable::Tensor const& b_scales, EpilogueArgs&&... epilogue_args) {
|
||||
STD_TORCH_CHECK(a.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
STD_TORCH_CHECK(b.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
|
||||
if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
return cutlass_gemm_sm100_fp8_batch_invariant_dispatch<
|
||||
cutlass::float_e4m3_t, cutlass::bfloat16_t, EnableBias>(
|
||||
out, a, b, a_scales, b_scales,
|
||||
std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
} else {
|
||||
STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
|
||||
return cutlass_gemm_sm100_fp8_batch_invariant_dispatch<
|
||||
cutlass::float_e4m3_t, cutlass::half_t, EnableBias>(
|
||||
out, a, b, a_scales, b_scales,
|
||||
std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
#include "scaled_mm_kernels.hpp"
|
||||
#include "scaled_mm_sm120_fp8_dispatch.cuh"
|
||||
#include "core/batch_invariant.hpp"
|
||||
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
|
||||
|
||||
namespace vllm {
|
||||
@@ -14,9 +15,17 @@ void cutlass_scaled_mm_sm120_fp8(
|
||||
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
|
||||
"currently bias dtype must match output dtype ",
|
||||
out.scalar_type());
|
||||
if (vllm_is_batch_invariant()) {
|
||||
return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue<
|
||||
c3x::ScaledEpilogueBias>(out, a, b, a_scales, b_scales, *bias);
|
||||
}
|
||||
return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogueBias>(
|
||||
out, a, b, a_scales, b_scales, *bias);
|
||||
} else {
|
||||
if (vllm_is_batch_invariant()) {
|
||||
return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue<
|
||||
c3x::ScaledEpilogue>(out, a, b, a_scales, b_scales);
|
||||
}
|
||||
return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogue>(
|
||||
out, a, b, a_scales, b_scales);
|
||||
}
|
||||
|
||||
@@ -179,6 +179,26 @@ inline void cutlass_gemm_sm120_fp8_dispatch(torch::stable::Tensor& out,
|
||||
out, a, b, std::forward<EpilogueArgs>(args)...);
|
||||
}
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename, typename> typename Epilogue,
|
||||
typename... EpilogueArgs>
|
||||
inline void cutlass_gemm_sm120_fp8_batch_invariant_dispatch(
|
||||
torch::stable::Tensor& out, torch::stable::Tensor const& a,
|
||||
torch::stable::Tensor const& b, EpilogueArgs&&... args) {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
STD_TORCH_CHECK(a.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
STD_TORCH_CHECK(b.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
|
||||
using Cutlass3xGemmM64 =
|
||||
typename sm120_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
|
||||
|
||||
// keep the CUTLASS config independent of M for batch invariance
|
||||
return cutlass_gemm_caller<Cutlass3xGemmM64>(
|
||||
out, a, b, std::forward<EpilogueArgs>(args)...);
|
||||
}
|
||||
|
||||
template <template <typename, typename, typename> typename Epilogue,
|
||||
typename... EpilogueArgs>
|
||||
void cutlass_scaled_mm_sm120_fp8_epilogue(torch::stable::Tensor& out,
|
||||
@@ -202,4 +222,26 @@ void cutlass_scaled_mm_sm120_fp8_epilogue(torch::stable::Tensor& out,
|
||||
}
|
||||
}
|
||||
|
||||
template <template <typename, typename, typename> typename Epilogue,
|
||||
typename... EpilogueArgs>
|
||||
void cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue(
|
||||
torch::stable::Tensor& out, torch::stable::Tensor const& a,
|
||||
torch::stable::Tensor const& b, EpilogueArgs&&... epilogue_args) {
|
||||
STD_TORCH_CHECK(a.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
STD_TORCH_CHECK(b.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
|
||||
if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
return cutlass_gemm_sm120_fp8_batch_invariant_dispatch<
|
||||
cutlass::float_e4m3_t, cutlass::bfloat16_t, Epilogue>(
|
||||
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
} else {
|
||||
STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
|
||||
return cutlass_gemm_sm120_fp8_batch_invariant_dispatch<
|
||||
cutlass::float_e4m3_t, cutlass::half_t, Epilogue>(
|
||||
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
#include "scaled_mm_kernels.hpp"
|
||||
#include "scaled_mm_sm90_fp8_dispatch.cuh"
|
||||
#include "core/batch_invariant.hpp"
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -13,9 +14,17 @@ void cutlass_scaled_mm_sm90_fp8(
|
||||
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
|
||||
"currently bias dtype must match output dtype ",
|
||||
out.scalar_type());
|
||||
if (vllm_is_batch_invariant()) {
|
||||
return cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue<true>(
|
||||
out, a, b, a_scales, b_scales, *bias);
|
||||
}
|
||||
return cutlass_scaled_mm_sm90_fp8_epilogue<true>(out, a, b, a_scales,
|
||||
b_scales, *bias);
|
||||
} else {
|
||||
if (vllm_is_batch_invariant()) {
|
||||
return cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue<false>(
|
||||
out, a, b, a_scales, b_scales);
|
||||
}
|
||||
return cutlass_scaled_mm_sm90_fp8_epilogue<false>(out, a, b, a_scales,
|
||||
b_scales);
|
||||
}
|
||||
|
||||
@@ -349,6 +349,35 @@ inline void cutlass_gemm_sm90_fp8_dispatch(
|
||||
}
|
||||
}
|
||||
|
||||
template <typename InType, typename OutType, bool EnableBias,
|
||||
typename... EpilogueArgs>
|
||||
inline void cutlass_gemm_sm90_fp8_batch_invariant_dispatch(
|
||||
torch::stable::Tensor& out, torch::stable::Tensor const& a,
|
||||
torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
|
||||
torch::stable::Tensor const& b_scales, EpilogueArgs&&... args) {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
STD_TORCH_CHECK(a.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
STD_TORCH_CHECK(b.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
|
||||
using Cutlass3xGemmM64_N1280 =
|
||||
typename sm90_fp8_config_M64_N1280<InType, OutType,
|
||||
EnableBias>::Cutlass3xGemm;
|
||||
using Cutlass3xGemmM64_N8192 =
|
||||
typename sm90_fp8_config_M64_N8192<InType, OutType,
|
||||
EnableBias>::Cutlass3xGemm;
|
||||
|
||||
// keep the CUTLASS config independent of M for batch invariance
|
||||
uint32_t const n = b.size(1);
|
||||
if (n <= 1280) {
|
||||
return cutlass_gemm_caller_sm90_fp8<Cutlass3xGemmM64_N1280>(
|
||||
out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
|
||||
}
|
||||
return cutlass_gemm_caller_sm90_fp8<Cutlass3xGemmM64_N8192>(
|
||||
out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
|
||||
}
|
||||
|
||||
template <bool EnableBias, typename... EpilogueArgs>
|
||||
void cutlass_scaled_mm_sm90_fp8_epilogue(torch::stable::Tensor& out,
|
||||
torch::stable::Tensor const& a,
|
||||
@@ -375,4 +404,28 @@ void cutlass_scaled_mm_sm90_fp8_epilogue(torch::stable::Tensor& out,
|
||||
}
|
||||
}
|
||||
|
||||
template <bool EnableBias, typename... EpilogueArgs>
|
||||
void cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue(
|
||||
torch::stable::Tensor& out, torch::stable::Tensor const& a,
|
||||
torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
|
||||
torch::stable::Tensor const& b_scales, EpilogueArgs&&... epilogue_args) {
|
||||
STD_TORCH_CHECK(a.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
STD_TORCH_CHECK(b.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
|
||||
if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
return cutlass_gemm_sm90_fp8_batch_invariant_dispatch<
|
||||
cutlass::float_e4m3_t, cutlass::bfloat16_t, EnableBias>(
|
||||
out, a, b, a_scales, b_scales,
|
||||
std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
} else {
|
||||
STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
|
||||
return cutlass_gemm_sm90_fp8_batch_invariant_dispatch<
|
||||
cutlass::float_e4m3_t, cutlass::half_t, EnableBias>(
|
||||
out, a, b, a_scales, b_scales,
|
||||
std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
#include "scaled_mm_c2x_sm89_fp8_dispatch.cuh"
|
||||
#include "scaled_mm_c2x_sm89_int8_dispatch.cuh"
|
||||
|
||||
#include "core/batch_invariant.hpp"
|
||||
#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp"
|
||||
|
||||
using namespace vllm;
|
||||
@@ -191,9 +192,17 @@ void cutlass_scaled_mm_sm89(torch::stable::Tensor& out,
|
||||
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
|
||||
"currently bias dtype must match output dtype ",
|
||||
out.scalar_type());
|
||||
if (vllm_is_batch_invariant()) {
|
||||
return cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue<
|
||||
c2x::ScaledEpilogueBias>(out, a, b, a_scales, b_scales, *bias);
|
||||
}
|
||||
return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBias>(
|
||||
out, a, b, a_scales, b_scales, *bias);
|
||||
} else {
|
||||
if (vllm_is_batch_invariant()) {
|
||||
return cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue<
|
||||
c2x::ScaledEpilogue>(out, a, b, a_scales, b_scales);
|
||||
}
|
||||
return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogue>(
|
||||
out, a, b, a_scales, b_scales);
|
||||
}
|
||||
|
||||
@@ -381,4 +381,43 @@ inline void cutlass_gemm_sm89_fp8_dispatch(torch::stable::Tensor& out,
|
||||
}
|
||||
}
|
||||
|
||||
template <typename InType, typename OutType,
|
||||
template <typename, typename> typename Epilogue,
|
||||
typename... EpilogueArgs>
|
||||
inline void cutlass_gemm_sm89_fp8_batch_invariant_dispatch(
|
||||
torch::stable::Tensor& out, torch::stable::Tensor const& a,
|
||||
torch::stable::Tensor const& b, EpilogueArgs&&... args) {
|
||||
static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
|
||||
STD_TORCH_CHECK(a.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
STD_TORCH_CHECK(b.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
|
||||
// keep the CUTLASS config independent of M for batch invariance
|
||||
return sm89_fp8_config_M64::dispatch<InType, OutType, Epilogue>(
|
||||
out, a, b, std::forward<EpilogueArgs>(args)...);
|
||||
}
|
||||
|
||||
template <template <typename, typename> typename Epilogue,
|
||||
typename... EpilogueArgs>
|
||||
void cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue(
|
||||
torch::stable::Tensor& out, torch::stable::Tensor const& a,
|
||||
torch::stable::Tensor const& b, EpilogueArgs&&... epilogue_args) {
|
||||
STD_TORCH_CHECK(a.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
STD_TORCH_CHECK(b.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Float8_e4m3fn);
|
||||
|
||||
if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
return cutlass_gemm_sm89_fp8_batch_invariant_dispatch<
|
||||
cutlass::float_e4m3_t, cutlass::bfloat16_t, Epilogue>(
|
||||
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
} else {
|
||||
STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
|
||||
return cutlass_gemm_sm89_fp8_batch_invariant_dispatch<
|
||||
cutlass::float_e4m3_t, cutlass::half_t, Epilogue>(
|
||||
out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
|
||||
@@ -156,17 +156,6 @@ inline int GetGroupsPerBlock(int64_t num_groups) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Largest divisor of padded_groups_per_row that is <= 16. ry = 16 / kx.
|
||||
inline int GetGroupsPerBlockX(int64_t padded_groups_per_row) {
|
||||
if (padded_groups_per_row % 16 == 0) {
|
||||
return 16;
|
||||
}
|
||||
if (padded_groups_per_row % 8 == 0) {
|
||||
return 8;
|
||||
}
|
||||
return 4;
|
||||
}
|
||||
|
||||
void per_token_group_quant_8bit(const torch::stable::Tensor& input,
|
||||
torch::stable::Tensor& output_q,
|
||||
torch::stable::Tensor& output_s,
|
||||
@@ -258,11 +247,11 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
|
||||
//
|
||||
// Constraints: GROUP_SIZE % (THREADS_PER_GROUP * VEC_SIZE) == 0; for
|
||||
// THREADS_PER_GROUP=8 and bf16/fp16 (VEC_SIZE=16), this means GROUP_SIZE=128.
|
||||
template <typename T, typename DST_DTYPE, int GROUP_SIZE, int kGroupsPerBlockX,
|
||||
int kRowsPerBlock>
|
||||
template <typename T, typename DST_DTYPE, int GROUP_SIZE>
|
||||
__global__ void per_token_group_quant_8bit_packed_register_kernel(
|
||||
const T* __restrict__ input, void* __restrict__ output_q,
|
||||
unsigned int* __restrict__ output_s_packed, const int padded_groups_per_row,
|
||||
unsigned int* __restrict__ output_s_packed, const int64_t num_groups_padded,
|
||||
const int groups_per_block, const int padded_groups_per_row,
|
||||
const int groups_per_row, const int mn, const int output_q_mn_extent,
|
||||
const int tma_aligned_mn, const int64_t num_scale_elems, const float eps,
|
||||
const float min_8bit, const float max_8bit) {
|
||||
@@ -271,25 +260,27 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
|
||||
constexpr int VEC_SIZE = 32 / sizeof(T); // 16 for bf16/fp16
|
||||
static_assert(GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE,
|
||||
"GROUP_SIZE must equal THREADS_PER_GROUP * VEC_SIZE");
|
||||
// Each group's 8 threads must live in a single warp octet so the
|
||||
// 0xffu << (threadIdx.x & 24u) shuffle mask selects exactly the lanes
|
||||
// that share a group. Requires 32 % THREADS_PER_GROUP == 0 and the host
|
||||
// to launch num_threads as a multiple of THREADS_PER_GROUP (which it does
|
||||
// via num_threads = groups_per_block * THREADS_PER_GROUP).
|
||||
static_assert(32 % THREADS_PER_GROUP == 0,
|
||||
"THREADS_PER_GROUP must divide warp size for the shuffle "
|
||||
"mask to be valid");
|
||||
static_assert(
|
||||
kGroupsPerBlockX > 0 && (kGroupsPerBlockX & (kGroupsPerBlockX - 1)) == 0,
|
||||
"kGroupsPerBlockX must be a positive power of 2");
|
||||
static_assert(kRowsPerBlock > 0, "kRowsPerBlock must be positive");
|
||||
|
||||
const int local_group_id = threadIdx.x / THREADS_PER_GROUP;
|
||||
const int lane_id = threadIdx.x % THREADS_PER_GROUP;
|
||||
|
||||
const int sf_k_local = local_group_id % kGroupsPerBlockX;
|
||||
const int row_local = local_group_id / kGroupsPerBlockX;
|
||||
const int sf_k_idx = blockIdx.x * kGroupsPerBlockX + sf_k_local;
|
||||
const int mn_idx = blockIdx.y * kRowsPerBlock + row_local;
|
||||
|
||||
if (mn_idx >= tma_aligned_mn) {
|
||||
const int64_t block_group_id = blockIdx.x * groups_per_block;
|
||||
const int64_t global_group_id = block_group_id + local_group_id;
|
||||
if (global_group_id >= num_groups_padded) {
|
||||
return;
|
||||
}
|
||||
|
||||
const int sf_k_idx =
|
||||
static_cast<int>(global_group_id % padded_groups_per_row);
|
||||
const int mn_idx = static_cast<int>(global_group_id / padded_groups_per_row);
|
||||
const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
|
||||
|
||||
// Load 16 input elements (32 B) into registers as two adjacent uint4
|
||||
@@ -452,53 +443,34 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
|
||||
|
||||
constexpr int THREADS_PER_GROUP = 8;
|
||||
const int64_t padded_groups_per_row = k_num_packed_sfk * 4;
|
||||
const int64_t num_groups_padded = tma_aligned_mn * padded_groups_per_row;
|
||||
const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn;
|
||||
|
||||
STD_TORCH_CHECK(padded_groups_per_row % 4 == 0,
|
||||
"padded_groups_per_row=", padded_groups_per_row,
|
||||
" is not a multiple of 4.");
|
||||
const int kx = GetGroupsPerBlockX(padded_groups_per_row);
|
||||
const int ry = 16 / kx;
|
||||
const int64_t blocks_x = padded_groups_per_row / kx;
|
||||
const int64_t blocks_y = (tma_aligned_mn + ry - 1) / ry;
|
||||
const int num_threads = (kx * ry) * THREADS_PER_GROUP;
|
||||
// CUDA caps grid.x and grid.y at 2^31 - 1; guard against pathological inputs.
|
||||
STD_TORCH_CHECK(blocks_x <= static_cast<int64_t>(INT32_MAX) &&
|
||||
blocks_y <= static_cast<int64_t>(INT32_MAX),
|
||||
"per_token_group_quant_8bit_packed grid too large: (",
|
||||
blocks_x, ", ", blocks_y, ").");
|
||||
const int groups_per_block = GetGroupsPerBlock(num_groups_padded);
|
||||
|
||||
auto dst_type = output_q.scalar_type();
|
||||
const int64_t num_blocks = num_groups_padded / groups_per_block;
|
||||
const int num_threads = groups_per_block * THREADS_PER_GROUP;
|
||||
// CUDA caps grid.x at 2^31 - 1; this fits any realistic shape but guard
|
||||
// against pathological inputs.
|
||||
STD_TORCH_CHECK(num_blocks <= static_cast<int64_t>(INT32_MAX),
|
||||
"per_token_group_quant_8bit_packed grid too large: ",
|
||||
num_blocks, " blocks (max ", INT32_MAX, ").");
|
||||
|
||||
#define LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, KX, RY) \
|
||||
do { \
|
||||
dim3 grid(static_cast<unsigned int>(blocks_x), \
|
||||
static_cast<unsigned int>(blocks_y)); \
|
||||
dim3 block(num_threads); \
|
||||
per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128, KX, \
|
||||
RY> \
|
||||
<<<grid, block, 0, stream>>>( \
|
||||
static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
|
||||
reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
|
||||
static_cast<int>(padded_groups_per_row), \
|
||||
static_cast<int>(groups_per_row), static_cast<int>(mn), \
|
||||
static_cast<int>(output_q_mn_extent), \
|
||||
static_cast<int>(tma_aligned_mn), num_scale_elems, \
|
||||
static_cast<float>(eps), static_cast<float>(min_8bit), \
|
||||
static_cast<float>(max_8bit)); \
|
||||
} while (0)
|
||||
|
||||
#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
|
||||
do { \
|
||||
if (kx == 16) { \
|
||||
LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 16, 1); \
|
||||
} else if (kx == 8) { \
|
||||
LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 8, 2); \
|
||||
} else if (kx == 4) { \
|
||||
LAUNCH_REG_KERNEL_INST(T, DST_DTYPE, 4, 4); \
|
||||
} else { \
|
||||
STD_TORCH_CHECK(false, "Unsupported kx value ", kx); \
|
||||
} \
|
||||
#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
|
||||
do { \
|
||||
dim3 grid(static_cast<unsigned int>(num_blocks)); \
|
||||
dim3 block(num_threads); \
|
||||
per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128> \
|
||||
<<<grid, block, 0, stream>>>( \
|
||||
static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
|
||||
reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
|
||||
num_groups_padded, groups_per_block, \
|
||||
static_cast<int>(padded_groups_per_row), \
|
||||
static_cast<int>(groups_per_row), static_cast<int>(mn), \
|
||||
static_cast<int>(output_q_mn_extent), \
|
||||
static_cast<int>(tma_aligned_mn), num_scale_elems, \
|
||||
static_cast<float>(eps), static_cast<float>(min_8bit), \
|
||||
static_cast<float>(max_8bit)); \
|
||||
} while (0)
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
@@ -516,7 +488,6 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
|
||||
}));
|
||||
|
||||
#undef LAUNCH_REG_KERNEL
|
||||
#undef LAUNCH_REG_KERNEL_INST
|
||||
}
|
||||
|
||||
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
|
||||
|
||||
@@ -0,0 +1,204 @@
|
||||
#include <Python.h>
|
||||
|
||||
extern "C" {
|
||||
|
||||
#include <stdbool.h>
|
||||
#include <time.h>
|
||||
|
||||
#if defined(__i386__) || defined(__x86_64__)
|
||||
#include <cpuid.h>
|
||||
#include <mwaitxintrin.h>
|
||||
#endif
|
||||
|
||||
#if defined(CLOCK_MONOTONIC_RAW)
|
||||
#define TIMEOUT_CLOCK CLOCK_MONOTONIC_RAW
|
||||
#else
|
||||
#define TIMEOUT_CLOCK CLOCK_MONOTONIC
|
||||
#endif
|
||||
|
||||
#define CPU_SUPPORT_NONE 0
|
||||
#define CPU_SUPPORT_MONITORX 1
|
||||
|
||||
#define MWAITX_DEFAULT_TIMEOUT_CYCLES 1000000
|
||||
|
||||
typedef struct {
|
||||
unsigned int cpu_support;
|
||||
unsigned int max_monitor_line_size;
|
||||
} spinloop_state_t;
|
||||
|
||||
static void determine_cpu_support(spinloop_state_t* state) {
|
||||
state->cpu_support = CPU_SUPPORT_NONE;
|
||||
state->max_monitor_line_size = 0;
|
||||
|
||||
#if defined(__i386__) || defined(__x86_64__)
|
||||
unsigned int eax, ebx, ecx, edx;
|
||||
if (__get_cpuid(0, &eax, &ebx, &ecx, &edx) == 1) {
|
||||
// AMD CPU (possible monitorx/mwaitx support)
|
||||
if (ebx == 0x68747541 && edx == 0x69746e65 && ecx == 0x444d4163) {
|
||||
if (__get_cpuid(0x80000000, &eax, &ebx, &ecx, &edx) == 1 &&
|
||||
eax >= 0x80000001 &&
|
||||
__get_cpuid(0x80000001, &eax, &ebx, &ecx, &edx) == 1) {
|
||||
if ((ecx & (1 << 29)) != 0) {
|
||||
state->cpu_support = CPU_SUPPORT_MONITORX;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (state->cpu_support == CPU_SUPPORT_MONITORX) {
|
||||
if (__get_cpuid(5, &eax, &ebx, &ecx, &edx) == 1) {
|
||||
state->max_monitor_line_size = ebx & 0xff;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
static PyObject* method_spinloop(PyObject* self, PyObject* args,
|
||||
PyObject* kwargs) {
|
||||
Py_buffer buffer;
|
||||
PyObject* callback;
|
||||
double timeout = 0.;
|
||||
|
||||
spinloop_state_t* state = (spinloop_state_t*)PyModule_GetState(self);
|
||||
if (state == NULL) {
|
||||
PyErr_SetString(PyExc_TypeError, "Failed to retrieve module state!");
|
||||
return NULL;
|
||||
}
|
||||
|
||||
static const char* keywords[] = {"buffer", "callback", "timeout", NULL};
|
||||
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*O|d", (char**)keywords,
|
||||
&buffer, &callback, &timeout)) {
|
||||
return NULL;
|
||||
}
|
||||
|
||||
if (!PyCallable_Check(callback)) {
|
||||
PyErr_SetString(PyExc_TypeError, "callback parameter must be callable!");
|
||||
PyBuffer_Release(&buffer);
|
||||
return NULL;
|
||||
}
|
||||
|
||||
struct timespec t_start;
|
||||
if (clock_gettime(TIMEOUT_CLOCK, &t_start) != 0) {
|
||||
PyErr_SetString(PyExc_RuntimeError, "clock_gettime() failed!");
|
||||
PyBuffer_Release(&buffer);
|
||||
return NULL;
|
||||
}
|
||||
|
||||
bool result = false;
|
||||
bool error = false;
|
||||
bool have_timeout = (timeout > 1e-9);
|
||||
unsigned int iteration = 0;
|
||||
const bool buffer_qualifies = (buffer.len <= state->max_monitor_line_size);
|
||||
|
||||
while (true) {
|
||||
PyObject* res = PyObject_CallNoArgs(callback);
|
||||
if (res == NULL) {
|
||||
error = true;
|
||||
break;
|
||||
}
|
||||
int ok = (res == Py_True);
|
||||
Py_DECREF(res);
|
||||
|
||||
if (ok) {
|
||||
result = true;
|
||||
break;
|
||||
}
|
||||
|
||||
// Check timeout at most every 16 iterations to avoid clock_gettime and
|
||||
// comparison cost
|
||||
if (have_timeout && (iteration & 15u) == 0) {
|
||||
struct timespec t_now;
|
||||
if (clock_gettime(TIMEOUT_CLOCK, &t_now) != 0) {
|
||||
PyErr_SetString(PyExc_RuntimeError, "clock_gettime() failed!");
|
||||
error = true;
|
||||
break;
|
||||
}
|
||||
|
||||
const double elapsed = (double)(t_now.tv_sec - t_start.tv_sec) +
|
||||
(t_now.tv_nsec - t_start.tv_nsec) * 1e-9;
|
||||
if (elapsed >= timeout) {
|
||||
result = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
++iteration;
|
||||
|
||||
#if defined(__i386__) || defined(__x86_64__)
|
||||
// monitorx + mwaitx with qualified buffer
|
||||
if (buffer_qualifies && state->cpu_support == CPU_SUPPORT_MONITORX) {
|
||||
_mm_monitorx(buffer.buf, 0, 0);
|
||||
|
||||
// Check once more in case the buffer has been modified while we were
|
||||
// arming the monitor hardware
|
||||
res = PyObject_CallNoArgs(callback);
|
||||
if (res == NULL) {
|
||||
error = true;
|
||||
break;
|
||||
}
|
||||
ok = (res == Py_True);
|
||||
Py_DECREF(res);
|
||||
|
||||
if (ok) {
|
||||
result = true;
|
||||
break;
|
||||
}
|
||||
|
||||
// Run mwaitx with enabled timeout (bit 1). The actual timeout value
|
||||
// is not very important, we just want to ensure we don't lock up
|
||||
// here for too long.
|
||||
Py_BEGIN_ALLOW_THREADS _mm_mwaitx((1 << 1), 0,
|
||||
MWAITX_DEFAULT_TIMEOUT_CYCLES);
|
||||
Py_END_ALLOW_THREADS
|
||||
}
|
||||
|
||||
// Fallback: Busy poll
|
||||
else {
|
||||
#endif
|
||||
// Give other threads a chance to be scheduled
|
||||
Py_BEGIN_ALLOW_THREADS
|
||||
#if defined(__i386__) || defined(__x86_64__)
|
||||
__builtin_ia32_pause();
|
||||
#elif defined(__aarch64__)
|
||||
__asm__ volatile("yield" :: : "memory");
|
||||
#endif
|
||||
Py_END_ALLOW_THREADS
|
||||
#if defined(__i386__) || defined(__x86_64__)
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
PyBuffer_Release(&buffer);
|
||||
|
||||
if (error) {
|
||||
return NULL;
|
||||
}
|
||||
|
||||
if (result) {
|
||||
Py_RETURN_TRUE;
|
||||
}
|
||||
|
||||
Py_RETURN_FALSE;
|
||||
}
|
||||
|
||||
static PyMethodDef spinloop_methods[] = {
|
||||
{"spinloop", (PyCFunction)method_spinloop, METH_VARARGS | METH_KEYWORDS,
|
||||
"Wait for store with callback"},
|
||||
{NULL, NULL, 0, NULL}};
|
||||
|
||||
static struct PyModuleDef spinloop_module = {
|
||||
PyModuleDef_HEAD_INIT, "spinloop",
|
||||
"Hardware-optimized spinloops for Python", sizeof(spinloop_state_t),
|
||||
spinloop_methods};
|
||||
|
||||
PyMODINIT_FUNC PyInit_spinloop(void) {
|
||||
PyObject* m = PyModule_Create(&spinloop_module);
|
||||
if (m != NULL) {
|
||||
spinloop_state_t* state = (spinloop_state_t*)PyModule_GetState(m);
|
||||
if (state != NULL) {
|
||||
determine_cpu_support(state);
|
||||
}
|
||||
}
|
||||
return m;
|
||||
}
|
||||
|
||||
} // extern "C"
|
||||
+4
-20
@@ -199,10 +199,7 @@ COPY requirements/cuda.txt requirements/cuda.txt
|
||||
COPY use_existing_torch.py use_existing_torch.py
|
||||
COPY pyproject.toml pyproject.toml
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
|
||||
sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' requirements/cuda.txt; \
|
||||
fi \
|
||||
&& if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
|
||||
if [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
|
||||
echo "Installing torch nightly..." \
|
||||
&& uv pip install --python /opt/venv/bin/python3 torch torchaudio torchvision --pre \
|
||||
--index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
|
||||
@@ -304,15 +301,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
python3 use_existing_torch.py --prefix; \
|
||||
fi
|
||||
|
||||
# Provision a bare interpreter for each CPython covered by `requires-python`
|
||||
# so DeepGEMM `_C` is built once per Python and bundled side-by-side in the
|
||||
# wheel; cmake reads DEEPGEMM_PYTHON_INTERPRETERS in deepgemm.cmake's
|
||||
# foreach loop. The matrix is derived from pyproject.toml.
|
||||
COPY tools/setup_deepgemm_pythons.sh tools/build_deepgemm_C.py tools/
|
||||
ENV DEEPGEMM_VENV_PREFIX=/opt/dgenv
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
tools/setup_deepgemm_pythons.sh > /tmp/dg_pythons.txt
|
||||
|
||||
# Build the vLLM wheel
|
||||
# if USE_SCCACHE is set, use sccache to speed up compilation
|
||||
# AWS credentials mounted at ~/.aws/credentials for sccache S3 auth (optional)
|
||||
@@ -340,7 +328,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
&& export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" \
|
||||
&& export VLLM_MAIN_CUDA_VERSION="${VLLM_MAIN_CUDA_VERSION}" \
|
||||
&& export VLLM_DOCKER_BUILD_CONTEXT=1 \
|
||||
&& export DEEPGEMM_PYTHON_INTERPRETERS=$(cat /tmp/dg_pythons.txt) \
|
||||
&& sccache --show-stats \
|
||||
&& python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \
|
||||
&& sccache --show-stats; \
|
||||
@@ -358,7 +345,6 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" && \
|
||||
export VLLM_PRECOMPILED_WHEEL_COMMIT="${VLLM_MERGE_BASE_COMMIT}" && \
|
||||
export VLLM_DOCKER_BUILD_CONTEXT=1 && \
|
||||
export DEEPGEMM_PYTHON_INTERPRETERS=$(cat /tmp/dg_pythons.txt) && \
|
||||
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
|
||||
fi
|
||||
|
||||
@@ -630,9 +616,6 @@ ARG PYTORCH_CUDA_INDEX_BASE_URL
|
||||
COPY requirements/common.txt /tmp/common.txt
|
||||
COPY requirements/cuda.txt /tmp/requirements-cuda.txt
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
if [ "$(echo $CUDA_VERSION | cut -d. -f1)" = "12" ]; then \
|
||||
sed -i 's/^nvidia-cutlass-dsl\[cu13\]>=/nvidia-cutlass-dsl>=/' /tmp/requirements-cuda.txt; \
|
||||
fi && \
|
||||
uv pip install --system -r /tmp/requirements-cuda.txt \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') && \
|
||||
rm /tmp/requirements-cuda.txt /tmp/common.txt
|
||||
@@ -846,6 +829,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
libcublas-dev-${CUDA_VERSION_DASH} \
|
||||
libcusolver-dev-${CUDA_VERSION_DASH}"; \
|
||||
if [ "$INSTALL_KV_CONNECTORS" = "true" ]; then \
|
||||
if [ "$CUDA_MAJOR" -ge 13 ]; then \
|
||||
uv pip install --system nixl-cu13; \
|
||||
fi; \
|
||||
uv pip install --system -r /tmp/kv_connectors.txt --no-build || ( \
|
||||
# if the above fails, install from source
|
||||
apt-get update -y && \
|
||||
@@ -855,8 +841,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# clean up -dev packages, keep runtime libraries
|
||||
rm -rf /var/lib/apt/lists/* \
|
||||
); \
|
||||
# Force-reinstall the matching CUDA wheel so the correct nixl_ep_cpp.so is installed.
|
||||
uv pip install --system --force-reinstall --no-deps nixl-cu${CUDA_MAJOR}; \
|
||||
fi
|
||||
|
||||
ENV VLLM_USAGE_SOURCE production-docker-image
|
||||
|
||||
@@ -39,7 +39,8 @@ th {
|
||||
| Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
|
||||
| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` |
|
||||
| Custom | ✅ | ✅ | Local file: `data.jsonl` |
|
||||
| Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
|
||||
| Custom Audio | ✅ | ✅ | Local file: `audio_data.jsonl` |
|
||||
| Custom Image | ✅ | ✅ | Local file: `image_data.jsonl` |
|
||||
|
||||
Legend:
|
||||
|
||||
@@ -142,7 +143,7 @@ Example output:  which do not require a "prompt" field; and multimodal audio-text chat models (e.g. Qwen2-Audio). Since these model types require different arguments at inference, we are giving two examples.
|
||||
|
||||
- **Example 1: Whisper**
|
||||
|
||||
Whisper is a dedicated ASR encoder-decoder model, so it uses `--backend openai-audio` and `--endpoint /v1/audio/transcriptions`.
|
||||
|
||||
```bash
|
||||
# start server
|
||||
vllm serve openai/whisper-tiny
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--model openai/whisper-tiny \
|
||||
--backend openai-audio \
|
||||
--endpoint /v1/audio/transcriptions \
|
||||
--dataset-name custom_audio \
|
||||
--dataset-path audio_data.jsonl \
|
||||
--no-oversample \
|
||||
--custom-output-len 256 \
|
||||
--save-result \
|
||||
--save-detailed \
|
||||
--result-filename whisper_bench.json
|
||||
```
|
||||
|
||||
- **Example 2: Qwen2-Audio**
|
||||
|
||||
Qwen2-Audio is a multimodal chat model that can do ASR and speech analysis, so it uses `--backend openai-chat`, and `--endpoint /v1/chat/completions`. It also requires `--enable-multimodal-chat` to enable multimodal chat transformation.
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--model Qwen/Qwen2-Audio-7B-Instruct \
|
||||
--backend openai-chat \
|
||||
--endpoint /v1/chat/completions \
|
||||
--dataset-name custom_audio \
|
||||
--dataset-path audio_data.jsonl \
|
||||
--no-oversample \
|
||||
--custom-output-len 256 \
|
||||
--enable-multimodal-chat \
|
||||
--save-result \
|
||||
--save-detailed \
|
||||
--result-filename qwen_bench.json
|
||||
```
|
||||
|
||||
#### Custom Image Dataset
|
||||
|
||||
If the image dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomImageDataset`. At inference time, use the option `--dataset-name custom_image`. Your data needs to be in the `.jsonl` format and needs to have "prompt" and "image_files" fields per entry, e.g., `image_data.jsonl`:
|
||||
|
||||
```json
|
||||
{"prompt": "How many animals are present in the given image?", "image_files": ["/path/to/image/folder/horsepony.jpg"]}
|
||||
@@ -193,8 +247,8 @@ vllm bench serve--save-result --save-detailed \
|
||||
--backend openai-chat \
|
||||
--model Qwen/Qwen2-VL-7B-Instruct \
|
||||
--endpoint /v1/chat/completions \
|
||||
--dataset-name custom_mm \
|
||||
--dataset-path <path-to-your-mm-data-jsonl> \
|
||||
--dataset-name custom_image \
|
||||
--dataset-path <path-to-your-image-data-jsonl> \
|
||||
--allowed-local-media-path /path/to/image/folder
|
||||
```
|
||||
|
||||
|
||||
@@ -125,13 +125,12 @@ Priority is **1 = highest** (tried first).
|
||||
| Priority | Backend |
|
||||
| -------- | ------- |
|
||||
| 1 | `FLASHINFER_MLA` |
|
||||
| 2 | `TOKENSPEED_MLA` |
|
||||
| 3 | `CUTLASS_MLA` |
|
||||
| 4 | `FLASH_ATTN_MLA` |
|
||||
| 5 | `FLASHMLA` |
|
||||
| 6 | `TRITON_MLA` |
|
||||
| 7 | `FLASHINFER_MLA_SPARSE`**\*** |
|
||||
| 8 | `FLASHMLA_SPARSE` |
|
||||
| 2 | `CUTLASS_MLA` |
|
||||
| 3 | `FLASH_ATTN_MLA` |
|
||||
| 4 | `FLASHMLA` |
|
||||
| 5 | `TRITON_MLA` |
|
||||
| 6 | `FLASHINFER_MLA_SPARSE`**\*** |
|
||||
| 7 | `FLASHMLA_SPARSE` |
|
||||
|
||||
**Ampere/Hopper (SM 8.x-9.x):**
|
||||
|
||||
@@ -203,7 +202,6 @@ hardware and configuration.
|
||||
| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | FA4 on SM100+, FA3 on SM90, FA2 otherwise |
|
||||
| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | DeepSeek R1 dims only |
|
||||
| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only |
|
||||
| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | DeepSeek R1 dims only |
|
||||
|
||||
> **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
|
||||
> On other GPUs, FlashAttention is used as the default.
|
||||
@@ -224,6 +222,5 @@ MLA decode backends are selected using the standard
|
||||
| `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`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `TOKENSPEED_MLA` | fp16, bf16 | `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `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 |
|
||||
|
||||
@@ -1,161 +0,0 @@
|
||||
# MooncakeStoreConnector Usage Guide
|
||||
|
||||
MooncakeStoreConnector is a KV cache connector that uses [MooncakeDistributedStore](https://github.com/kvcache-ai/Mooncake) as a shared KV cache pool. Unlike `MooncakeConnector` which does direct point-to-point KV transfer between prefiller and decoder, MooncakeStoreConnector enables KV cache offloading to an external distributed store, supporting:
|
||||
|
||||
- **CPU offloading**: Extend effective KV cache capacity by offloading to CPU memory via Mooncake's transfer engine.
|
||||
- **Prefix caching across instances**: Hash-based deduplication allows multiple vLLM instances to share cached KV blocks through the store.
|
||||
- **Single-node and multi-node deployment**: Works both as a standalone KV cache extension and in disaggregated prefill-decode setups.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
### Install Mooncake
|
||||
|
||||
Install mooncake through pip:
|
||||
|
||||
```bash
|
||||
uv pip install mooncake-transfer-engine
|
||||
```
|
||||
|
||||
Refer to the [Mooncake official repository](https://github.com/kvcache-ai/Mooncake) for more installation instructions and building from source.
|
||||
|
||||
### Start the Mooncake Master Server
|
||||
|
||||
The Mooncake master manages metadata and coordinates the distributed store. Start it before launching vLLM:
|
||||
|
||||
```bash
|
||||
mooncake_master --port 50051
|
||||
```
|
||||
|
||||
Default ports:
|
||||
|
||||
- RPC: 50051
|
||||
|
||||
Multiple vLLM instances can share the same master server.
|
||||
|
||||
### Configure Mooncake
|
||||
|
||||
Create a JSON configuration file (e.g., `mooncake_config.json`):
|
||||
|
||||
```json
|
||||
{
|
||||
"metadata_server": "P2PHANDSHAKE",
|
||||
"master_server_address": "127.0.0.1:50051",
|
||||
"global_segment_size": "80GB",
|
||||
"local_buffer_size": "4GB",
|
||||
"protocol": "rdma",
|
||||
"device_name": ""
|
||||
}
|
||||
```
|
||||
|
||||
- `protocol`: Use `"rdma"` for best performance. `"tcp"` works as a fallback.
|
||||
- `global_segment_size`: CPU memory contributed to the distributed pool (per GPU).
|
||||
- `local_buffer_size`: Private buffer for this node's own operations (per GPU).
|
||||
|
||||
Set the config path via environment variable:
|
||||
|
||||
```bash
|
||||
export MOONCAKE_CONFIG_PATH=/path/to/mooncake_config.json
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Single-Node KV Cache Offloading
|
||||
|
||||
Use MooncakeStoreConnector to offload KV cache to CPU memory, extending the effective cache size:
|
||||
|
||||
```bash
|
||||
MOONCAKE_CONFIG_PATH=mooncake_config.json \
|
||||
vllm serve meta-llama/Llama-3.1-8B-Instruct \
|
||||
--kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both"}'
|
||||
```
|
||||
|
||||
### Disaggregated Prefill-Decode (XpYd)
|
||||
|
||||
In disaggregated prefill-decode mode, use `MultiConnector` to combine `MooncakeConnector` (point-to-point KV transfer) with `MooncakeStoreConnector` (shared KV cache pool). This enables both direct P2P transfer between prefiller and decoder, and cross-instance prefix cache sharing via the distributed store.
|
||||
**Prefiller Node:**
|
||||
|
||||
```bash
|
||||
MOONCAKE_CONFIG_PATH=mooncake_config.json \
|
||||
VLLM_MOONCAKE_BOOTSTRAP_PORT=50052 \
|
||||
vllm serve meta-llama/Llama-3.1-8B-Instruct \
|
||||
--port 8100 \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "MultiConnector",
|
||||
"kv_role": "kv_producer",
|
||||
"kv_connector_extra_config": {
|
||||
"connectors": [
|
||||
{
|
||||
"kv_connector": "MooncakeConnector",
|
||||
"kv_role": "kv_producer"
|
||||
},
|
||||
{
|
||||
"kv_connector": "MooncakeStoreConnector",
|
||||
"kv_role": "kv_producer"
|
||||
}
|
||||
]
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
**Decoder Node:**
|
||||
|
||||
```bash
|
||||
MOONCAKE_CONFIG_PATH=mooncake_config.json \
|
||||
VLLM_MOONCAKE_BOOTSTRAP_PORT=50053 \
|
||||
vllm serve meta-llama/Llama-3.1-8B-Instruct \
|
||||
--port 8200 \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "MultiConnector",
|
||||
"kv_role": "kv_consumer",
|
||||
"kv_connector_extra_config": {
|
||||
"connectors": [
|
||||
{
|
||||
"kv_connector": "MooncakeConnector",
|
||||
"kv_role": "kv_consumer"
|
||||
},
|
||||
{
|
||||
"kv_connector": "MooncakeStoreConnector",
|
||||
"kv_role": "kv_consumer"
|
||||
}
|
||||
]
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
**Proxy:**
|
||||
|
||||
A disaggregation proxy is required to route requests between prefiller and decoder nodes. The proxy assigns `do_remote_prefill=True` / `do_remote_decode=True` to coordinate P2P transfer via `MooncakeConnector`. Refer to the [MooncakeConnector usage guide](mooncake_connector_usage.md) for proxy setup details.
|
||||
|
||||
## Environment Variables
|
||||
|
||||
| Variable | Description | Default |
|
||||
| --- | --- | --- |
|
||||
| `MOONCAKE_CONFIG_PATH` | Path to Mooncake JSON config file | (required) |
|
||||
| `VLLM_MOONCAKE_BOOTSTRAP_PORT` | Bootstrap port for MooncakeConnector P2P transfer (disagg mode only) | 8998 |
|
||||
|
||||
## KV Transfer Config
|
||||
|
||||
### KV Role Options
|
||||
|
||||
- **kv_producer**: For prefiller instances that store KV caches to the pool.
|
||||
- **kv_consumer**: For decoder instances that load KV caches from the pool.
|
||||
- **kv_both**: The instance both stores and loads KV caches. Use this for single-node CPU offloading.
|
||||
|
||||
### kv_connector_extra_config
|
||||
|
||||
- `load_async` (bool): Enable asynchronous loading for better compute-I/O overlap. Default: `true`.
|
||||
- `enable_cross_layers_blocks` (bool): Enable cross-layer block packing for reduced store operations. Default: `false`.
|
||||
- `discard_partial_chunks` (bool): Discard partial block chunks during store. Default: `true`.
|
||||
- `lookup_rpc_port` (int): Custom port for the ZMQ lookup RPC socket. Default: `0`.
|
||||
|
||||
## Notes
|
||||
|
||||
### Cross-DP Prefix Cache Hits
|
||||
|
||||
When running with data parallelism, set a fixed `PYTHONHASHSEED` so that block hashes are consistent across DP ranks:
|
||||
|
||||
```bash
|
||||
PYTHONHASHSEED=0 vllm serve ...
|
||||
```
|
||||
|
||||
Without this, identical prompts may produce different block hashes on different DP ranks, preventing cross-instance prefix cache hits.
|
||||
@@ -126,9 +126,15 @@ python tests/v1/kv_connector/nixl_integration/toy_proxy_server.py \
|
||||
- Set when prefiller and decoder are on different machines
|
||||
- Connection info is passed via KVTransferParams from prefiller to decoder for handshake
|
||||
|
||||
- `VLLM_NIXL_ABORT_REQUEST_TIMEOUT`: Timeout (in seconds) for automatically releasing the prefiller’s KV cache for a particular request. (Optional)
|
||||
- `kv_lease_duration` (via `kv_connector_extra_config`): Lease duration (in seconds) for the prefiller's KV cache blocks. (Optional)
|
||||
- Default: 30
|
||||
- When a prefill request finishes, its KV blocks are held for this duration waiting for the decoder to read them. While the request is queued on the decoder, periodic heartbeats automatically extend the lease. If neither a heartbeat nor a read notification arrives before the lease expires, the blocks are freed. The heartbeat interval and extension amount are derived automatically from this value.
|
||||
- Example: `--kv-transfer-config '{"kv_connector_extra_config": {"kv_lease_duration": 60}}'`
|
||||
|
||||
- `decoder_kv_blocks_ttl` (via `kv_connector_extra_config`): TTL (in seconds) for KV blocks cached on the decoder in bidirectional transfer mode. (Optional)
|
||||
- Default: 480
|
||||
- If a request is aborted and the decoder has not yet read the KV-cache blocks through the nixl channel, the prefill instance will release its KV-cache blocks after this timeout to avoid holding them indefinitely.
|
||||
- In bidirectional mode, the decoder caches KV blocks for multi-turn conversations. This TTL controls how long those blocks are held before being released. Unlike the prefiller lease, this TTL is not renewed via heartbeats.
|
||||
- Example: `--kv-transfer-config '{"kv_connector_extra_config": {"decoder_kv_blocks_ttl": 600}}'`
|
||||
|
||||
## Multi-Instance Setup
|
||||
|
||||
|
||||
@@ -3,9 +3,10 @@
|
||||
vLLM supports the following hardware platforms:
|
||||
|
||||
- [GPU](gpu.md)
|
||||
- [NVIDIA CUDA](gpu.md#nvidia-cuda)
|
||||
- [AMD ROCm](gpu.md#amd-rocm)
|
||||
- [Intel XPU](gpu.md#intel-xpu)
|
||||
- [NVIDIA CUDA](gpu.md)
|
||||
- [AMD ROCm](gpu.md)
|
||||
- [Intel XPU](gpu.md)
|
||||
- [Apple Silicon](gpu.md) (via [vLLM-Metal](https://github.com/vllm-project/vllm-metal))
|
||||
- [CPU](cpu.md)
|
||||
- [Intel/AMD x86](cpu.md#intelamd-x86)
|
||||
- [ARM AArch64](cpu.md#arm-aarch64)
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
<!-- markdownlint-disable MD041 -->
|
||||
--8<-- [start:installation]
|
||||
|
||||
For GPU-accelerated inference on Apple Silicon, use [vLLM-Metal](https://github.com/vllm-project/vllm-metal), a community-maintained hardware plugin that uses MLX as the compute backend and provides native GPU acceleration via Apple's Metal framework.
|
||||
|
||||
vLLM-Metal works with MLX-optimized models from the [mlx-community](https://huggingface.co/mlx-community) organization on Hugging Face, which provides quantized versions of popular models optimized for Apple Silicon.
|
||||
|
||||
!!! tip
|
||||
For installation and usage instructions, see the [Set up using vLLM-Metal](#set-up-using-vllm-metal) section below.
|
||||
|
||||
--8<-- [end:installation]
|
||||
--8<-- [start:requirements]
|
||||
|
||||
- OS: macOS Sonoma or later
|
||||
- Hardware: Apple Silicon
|
||||
- Metal support enabled
|
||||
|
||||
!!! note
|
||||
See the [Set up using vLLM-Metal](#set-up-using-vllm-metal) section below for installation instructions.
|
||||
|
||||
--8<-- [end:requirements]
|
||||
--8<-- [start:set-up-using-python]
|
||||
|
||||
## Set up using vLLM-Metal
|
||||
|
||||
vLLM-Metal is distributed as a separate package that provides native GPU acceleration on Apple Silicon.
|
||||
|
||||
To install vLLM-Metal, follow the installation instructions in the [vLLM-Metal documentation](https://github.com/vllm-project/vllm-metal#installation).
|
||||
|
||||
The installation will:
|
||||
|
||||
1. Set up the appropriate Python environment
|
||||
2. Install MLX and required dependencies
|
||||
3. Install the vLLM-Metal package
|
||||
|
||||
After installation, you can start using vLLM with Metal GPU acceleration.
|
||||
|
||||
!!! tip
|
||||
When using vLLM-Metal, use models from the [mlx-community](https://huggingface.co/mlx-community) on Hugging Face for best performance. These models are optimized for MLX and often include quantized versions (4-bit, 8-bit) that run efficiently on Apple Silicon.
|
||||
|
||||
Example model: `mlx-community/Qwen2.5-0.5B-Instruct-4bit`
|
||||
|
||||
### Using vLLM-Metal
|
||||
|
||||
After installation, vLLM-Metal provides an easy-to-use CLI for running an OpenAI-compatible API server:
|
||||
|
||||
```bash
|
||||
# Activate the vLLM-Metal environment
|
||||
source ~/.venv-vllm-metal/bin/activate
|
||||
|
||||
# Start the API server (specify your mlx-community model or it will use default)
|
||||
vllm serve
|
||||
```
|
||||
|
||||
Once the server is running, you have multiple options to interact with it:
|
||||
|
||||
#### Option 1: Interactive chat
|
||||
|
||||
Open a new terminal and start an interactive chat session:
|
||||
|
||||
```bash
|
||||
source ~/.venv-vllm-metal/bin/activate
|
||||
vllm chat
|
||||
```
|
||||
|
||||
#### Option 2: API requests with curl
|
||||
|
||||
```bash
|
||||
curl http://localhost:8000/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": "Hello!"}],
|
||||
"max_tokens": 50
|
||||
}'
|
||||
```
|
||||
|
||||
#### Option 3: Python with OpenAI SDK
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(
|
||||
base_url="http://localhost:8000/v1",
|
||||
api_key="dummy" # No auth required for local server
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="mlx-community/Qwen2.5-0.5B-Instruct-4bit",
|
||||
messages=[{"role": "user", "content": "Hello!"}]
|
||||
)
|
||||
|
||||
print(response.choices[0].message.content)
|
||||
```
|
||||
|
||||
For more details on the `vllm` CLI commands, see the [OpenAI-compatible server documentation](../../serving/openai_compatible_server.md).
|
||||
|
||||
--8<-- [end:set-up-using-python]
|
||||
--8<-- [start:pre-built-wheels]
|
||||
|
||||
vLLM-Metal is installed via the vLLM-Metal package. See the [Set up using vLLM-Metal](#set-up-using-vllm-metal) section above.
|
||||
|
||||
--8<-- [end:pre-built-wheels]
|
||||
--8<-- [start:build-wheel-from-source]
|
||||
|
||||
For build instructions from source, refer to the [vLLM-Metal documentation](https://github.com/vllm-project/vllm-metal#installation).
|
||||
|
||||
--8<-- [end:build-wheel-from-source]
|
||||
--8<-- [start:pre-built-images]
|
||||
|
||||
--8<-- [end:pre-built-images]
|
||||
--8<-- [start:build-image-from-source]
|
||||
|
||||
--8<-- [end:build-image-from-source]
|
||||
--8<-- [start:supported-features]
|
||||
|
||||
vLLM-Metal provides:
|
||||
|
||||
- Native GPU acceleration using Metal
|
||||
- MLX-based compute backend optimized for Apple Silicon
|
||||
- OpenAI-compatible API server
|
||||
- Support for popular model architectures
|
||||
|
||||
For specific feature support and limitations, refer to the [vLLM-Metal documentation](https://github.com/vllm-project/vllm-metal).
|
||||
|
||||
--8<-- [end:supported-features]
|
||||
@@ -18,6 +18,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:installation"
|
||||
|
||||
=== "Apple Silicon"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:installation"
|
||||
|
||||
## Requirements
|
||||
|
||||
- OS: Linux
|
||||
@@ -38,6 +42,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:requirements"
|
||||
|
||||
=== "Apple Silicon"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:requirements"
|
||||
|
||||
## Set up using Python
|
||||
|
||||
### Create a new Python environment
|
||||
@@ -56,6 +64,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:set-up-using-python"
|
||||
|
||||
=== "Apple Silicon"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:set-up-using-python"
|
||||
|
||||
### Pre-built wheels {#pre-built-wheels}
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
@@ -70,6 +82,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:pre-built-wheels"
|
||||
|
||||
=== "Apple Silicon"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:pre-built-wheels"
|
||||
|
||||
### Build wheel from source
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
@@ -84,6 +100,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:build-wheel-from-source"
|
||||
|
||||
=== "Apple Silicon"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:build-wheel-from-source"
|
||||
|
||||
## Set up using Docker
|
||||
|
||||
### Pre-built images
|
||||
@@ -102,6 +122,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:pre-built-images"
|
||||
|
||||
=== "Apple Silicon"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:pre-built-images"
|
||||
|
||||
--8<-- [end:pre-built-images]
|
||||
|
||||
### Build image from source
|
||||
@@ -120,6 +144,10 @@ vLLM is a Python library that supports the following GPU variants. Select your G
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:build-image-from-source"
|
||||
|
||||
=== "Apple Silicon"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:build-image-from-source"
|
||||
|
||||
--8<-- [end:build-image-from-source]
|
||||
|
||||
## Supported features
|
||||
@@ -135,3 +163,7 @@ vLLM is a Python library that supports the following GPU variants. Select your G
|
||||
=== "Intel XPU"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.xpu.inc.md:supported-features"
|
||||
|
||||
=== "Apple Silicon"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu.apple.inc.md:supported-features"
|
||||
|
||||
@@ -10,6 +10,9 @@ This guide will help you quickly get started with vLLM to perform:
|
||||
- OS: Linux
|
||||
- Python: 3.10 -- 3.13
|
||||
|
||||
!!! note
|
||||
vLLM also works on macOS with [vLLM-Metal](https://github.com/vllm-project/vllm-metal) for Apple Silicon GPU acceleration. See the [GPU installation guide](installation/gpu.md) and select the "Apple Silicon" tab.
|
||||
|
||||
## Installation
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
@@ -73,6 +76,18 @@ This guide will help you quickly get started with vLLM to perform:
|
||||
!!! note
|
||||
For more detailed instructions, including Docker, installing from source, and troubleshooting, please refer to the [vLLM on TPU documentation](https://docs.vllm.ai/projects/tpu/en/latest/).
|
||||
|
||||
=== "Apple Silicon (Mac)"
|
||||
|
||||
If you are using Apple Silicon Macs, you can use vLLM-Metal for GPU-accelerated inference via Apple's Metal framework.
|
||||
|
||||
Follow the installation instructions in the [vLLM-Metal documentation](https://github.com/vllm-project/vllm-metal#installation).
|
||||
|
||||
!!! note
|
||||
vLLM-Metal uses MLX instead of PyTorch as the compute backend and requires MLX-optimized models from the [mlx-community](https://huggingface.co/mlx-community) on Hugging Face.
|
||||
|
||||
!!! tip
|
||||
For more detailed instructions, please refer to the [GPU installation guide](installation/gpu.md) and select the "Apple Silicon" tab.
|
||||
|
||||
!!! note
|
||||
For more detail and non-CUDA platforms, please refer to the [installation guide](installation/README.md) for specific instructions on how to install vLLM.
|
||||
|
||||
|
||||
@@ -598,7 +598,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
|
||||
| `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
|
||||
| `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
|
||||
| `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B` | ✅︎ | ✅︎ |
|
||||
| `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`<sup>^</sup>, `allenai/MolmoWeb-8B`<sup>^</sup> | ✅︎ | ✅︎ |
|
||||
| `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
|
||||
| `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ |
|
||||
| `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
|
||||
@@ -666,6 +666,11 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
!!! note
|
||||
For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.
|
||||
|
||||
!!! note
|
||||
To use `allenai/MolmoWeb-4B` or `allenai/MolmoWeb-8B`, serve the checkpoint
|
||||
with the Molmo2 architecture and disable multimodal-prefix attention:
|
||||
`--hf-overrides '{"architectures": ["Molmo2ForConditionalGeneration"], "is_mm_prefix_lm": false}'`.
|
||||
|
||||
!!! note
|
||||
`Moondream3ForCausalLM` uses task-specific prompt templates for `query`
|
||||
and `caption`. The native `detect` and `point` skills require custom
|
||||
|
||||
@@ -333,6 +333,43 @@ Most cache paths default to subdirectories under a single root. Changing `VLLM_C
|
||||
- **Do not copy cache contents from untrusted sources.** If you distribute cache artifacts between environments, ensure they originate from a trusted build pipeline.
|
||||
- **Container deployments:** If mounting cache directories into containers, ensure the volume source is trusted.
|
||||
|
||||
## FIPS Compatibility
|
||||
|
||||
FIPS compliance depends on many factors, so a vLLM deployment is not automatically FIPS compliant. Recent changes have improved vLLM's *tolerance* of FIPS-enabled hosts — that is, avoiding crashes when non-approved algorithms are blocked — but tolerance is not the same as compliance. Whether a deployment satisfies FIPS requirements depends on the host operating system, the OpenSSL provider backing Python's `hashlib` and `ssl` modules, and which optional dependencies are installed.
|
||||
|
||||
### FIPS-relevant configuration
|
||||
|
||||
Operators running vLLM on FIPS-enabled hosts should select FIPS-approved algorithms via the following knobs:
|
||||
|
||||
- **Multimodal input hashing** — `VLLM_MM_HASHER_ALGORITHM` defaults to `blake3`, which is not FIPS-approved. Set it to `sha256` or `sha512` in FIPS-enabled environments.
|
||||
- **Prefix-cache hashing** — set `--prefix-caching-hash-algo` (config field `prefix_caching_hash_algo`) to `sha256` or `sha256_cbor`. The `xxhash` and `xxhash_cbor` options are not FIPS-approved.
|
||||
- **TLS ciphers** — use `--ssl-ciphers` to restrict the API server's TLS handshake to FIPS-approved cipher suites that match your environment's policy.
|
||||
|
||||
### Automatic fallback for non-security MD5 use
|
||||
|
||||
vLLM uses MD5 in a few places to derive non-security cache keys (for example, configuration hashes). These call sites pass `usedforsecurity=False` and additionally fall back to SHA-256 when the underlying OpenSSL provider refuses MD5 outright (see `safe_hash()` in `vllm/utils/hashing.py`). No user action is required; this behavior is documented so that auditors and security reviewers can identify the MD5 references and understand their purpose.
|
||||
|
||||
### Dependencies that provide non-FIPS hash implementations
|
||||
|
||||
Some dependencies expose hash implementations that are not FIPS-approved. vLLM only invokes them when the corresponding algorithm is selected, but operators with strict cryptographic controls may want to ensure the code paths are not exercised — and, where policy requires, that the packages themselves are absent:
|
||||
|
||||
- `blake3` — currently listed in `requirements/common.txt`, so a standard install pulls it in. It is imported lazily and only used when `VLLM_MM_HASHER_ALGORITHM=blake3` (the default). Setting `VLLM_MM_HASHER_ALGORITHM` to `sha256` or `sha512` is sufficient to keep the non-FIPS code path dormant. If your policy additionally forbids the package being present, uninstall it after `pip install` (`pip uninstall blake3`); vLLM will continue to function as long as `VLLM_MM_HASHER_ALGORITHM` is set to a non-blake3 value.
|
||||
- `xxhash` — a true optional dependency (not in `requirements/common.txt`). It is only imported when an `xxhash`-based prefix-cache algorithm is selected. Leave it uninstalled and select a `sha256`-based prefix-cache algorithm.
|
||||
|
||||
### Beyond hashing: other FIPS considerations
|
||||
|
||||
Hashing is the area where vLLM has explicit FIPS-aware code, but a FIPS-compliant deployment depends on several factors that sit outside vLLM itself. Operators should evaluate the following with their platform and security teams:
|
||||
|
||||
- **Host crypto provider.** Python's `hashlib` and `ssl` modules are FIPS-aware only when Python is linked against a FIPS-validated OpenSSL (or equivalent) provider supplied by the host OS. vLLM inherits whatever provider the host configures — it does not bundle one.
|
||||
- **API server TLS.** TLS termination for the OpenAI-compatible API server uses the host's OpenSSL via Python's `ssl` module. Restrict the cipher suite with `--ssl-ciphers` to match your environment's FIPS policy, and ensure server certificates are issued with FIPS-approved algorithms and key sizes.
|
||||
- **Outbound HTTPS.** Model and asset downloads (for example, via `huggingface_hub`) use the same host TLS stack. The same provider/cipher considerations apply.
|
||||
- **Inter-node communication is unencrypted by default.** As described in [Inter-Node Communication](#inter-node-communication), PyTorch Distributed, KV-cache transfer, and data-parallel channels do not encrypt traffic. FIPS environments that require FIPS-approved cryptography for data in transit must provide that protection externally — for example, via an mTLS sidecar or IPsec terminated by a FIPS-validated module — since vLLM's internal channels cannot satisfy the requirement on their own. Network isolation alone is not cryptography and does not meet a "FIPS-approved cryptography for data in transit" requirement, though it remains a useful defense-in-depth measure.
|
||||
- **Dependencies that bundle their own OpenSSL.** Some Python wheels statically link OpenSSL builds that fail the kernel FIPS self-test on FIPS-enabled hosts (`FATAL FIPS SELFTEST FAILURE`). `opencv-python-headless` is a known example; other manylinux wheels may behave similarly. Audit your installed wheels for bundled crypto libraries when troubleshooting FIPS startup failures.
|
||||
- **Accelerator and ML libraries.** PyTorch, CUDA, cuDNN, NCCL, and similar components have their own crypto and FIPS posture independent of vLLM. NVIDIA publishes FIPS-validated builds for some libraries; vLLM does not pin to those builds, so selecting and validating them is the operator's responsibility.
|
||||
- **What is *not* a FIPS concern in vLLM.** Random number generation used for token sampling (Python/NumPy/PyTorch RNGs) is not a cryptographic use and is out of scope for FIPS. Pickled cache artifacts are a separate security concern covered under [Cache Directory Security](#cache-directory-security).
|
||||
|
||||
In short: the configuration knobs above let vLLM avoid non-approved algorithms, and the automatic fallbacks let it run without crashing on FIPS-enabled hosts. End-to-end FIPS compliance, however, is a property of the full deployment — host OS, crypto provider, transitive dependencies, and network architecture — not of vLLM alone.
|
||||
|
||||
## Reporting Security Vulnerabilities
|
||||
|
||||
If you believe you have found a security vulnerability in vLLM, please report it following the project's security policy. For more information on how to report security issues and the project's security policy, please see the [vLLM Security Policy](https://github.com/vllm-project/vllm/blob/main/SECURITY.md).
|
||||
|
||||
+2
-1
@@ -123,7 +123,8 @@ python = "./.venv"
|
||||
extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizers_/*",
|
||||
"benchmarks/sonnet.txt", "tests/lora/data/*", "build/*",
|
||||
"examples/pooling/token_embed/*", "tests/models/language/pooling/*",
|
||||
"vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", "tests/entrypoints/openai/speech_to_text/test_transcription_validation.py",
|
||||
"vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*",
|
||||
"tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py",
|
||||
"docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html",
|
||||
"tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*"]
|
||||
ignore-hidden = false
|
||||
|
||||
@@ -21,8 +21,5 @@ nvidia-cudnn-frontend>=1.13.0,<1.19.0
|
||||
fastsafetensors >= 0.2.2
|
||||
|
||||
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
|
||||
nvidia-cutlass-dsl==4.4.2
|
||||
nvidia-cutlass-dsl>=4.4.2
|
||||
quack-kernels>=0.3.3
|
||||
|
||||
# Tokenspeed_MLA for faster mla with spec decode
|
||||
tokenspeed-mla==0.1.2
|
||||
@@ -1,3 +1,5 @@
|
||||
lmcache >= 0.3.9
|
||||
nixl >= 1.1.0 # Required for disaggregated prefill
|
||||
nixl[cu13] >= 0.7.1, <= 0.10.1 # Required for disaggregated prefill
|
||||
nixl-cu12 >= 0.7.1, <= 0.10.1
|
||||
nixl-cu13 >= 0.7.1, <= 0.10.1
|
||||
mooncake-transfer-engine >= 0.3.8
|
||||
|
||||
@@ -686,6 +686,7 @@ class precompiled_wheel_utils:
|
||||
"vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
|
||||
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
|
||||
"vllm/cumem_allocator.abi3.so",
|
||||
"vllm/spinloop.abi3.so",
|
||||
# ROCm-specific libraries
|
||||
"vllm/_rocm_C.abi3.so",
|
||||
]
|
||||
@@ -969,9 +970,6 @@ def get_requirements() -> list[str]:
|
||||
# vllm-flash-attn is built only for CUDA 12.x.
|
||||
# Skip for other versions.
|
||||
continue
|
||||
if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12":
|
||||
# [cu13] extra is the default; strip it on CUDA 12 builds.
|
||||
req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl")
|
||||
modified_requirements.append(req)
|
||||
requirements = modified_requirements
|
||||
elif _is_hip():
|
||||
@@ -996,6 +994,8 @@ if _is_cuda() or _is_hip():
|
||||
# copying the relevant .py files from the source repository.
|
||||
ext_modules.append(CMakeExtension(name="vllm.triton_kernels", optional=True))
|
||||
|
||||
ext_modules.append(CMakeExtension(name="vllm.spinloop"))
|
||||
|
||||
if _is_hip():
|
||||
ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
|
||||
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Unit tests for the DoubleQuant fan-out variants registered by
|
||||
``RocmAiterRMSNormQuantFusionPass``.
|
||||
|
||||
Both variants target a 1-to-2 fan-out where one ``rms_norm`` output feeds
|
||||
two distinct ``rocm_aiter_group_fp8_quant`` consumers and rewrite it into
|
||||
two independent fused ``rms_norm + group_fp8_quant`` ops:
|
||||
|
||||
* ``DoubleAiterRMSFp8GroupQuantPattern`` matches the un-viewed shape
|
||||
(e.g. Kimi-K2.5 / DSR1).
|
||||
* ``DoubleAiterRMSFp8GroupQuantViewPattern`` (this PR) is the view-tolerant
|
||||
sibling that additionally matches the
|
||||
``rms_norm -> view -> group_fp8_quant`` shape that DSv3.2's MLA indexer
|
||||
q_c norm exposes through ``Fp8BlockScaledMMLinearKernel.apply_weights``'s
|
||||
2D-flatten boilerplate.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm.config
|
||||
from tests.compile.backend import TestBackend
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
|
||||
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
|
||||
from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
|
||||
from vllm.config import (
|
||||
CompilationConfig,
|
||||
CompilationMode,
|
||||
ModelConfig,
|
||||
PassConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
|
||||
EPS = 1e-5
|
||||
HIDDEN_SIZE = 256
|
||||
GROUP_SIZE = 128
|
||||
|
||||
|
||||
class _NoViewDoubleQuantModel(torch.nn.Module):
|
||||
"""``rms_norm -> 2x group_fp8_quant`` fan-out (Kimi-K2.5 / DSR1 shape)."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.weight = torch.nn.Parameter(torch.ones(HIDDEN_SIZE, dtype=torch.bfloat16))
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
# avoid graph input being a direct arg to a matched pattern node
|
||||
x = torch.relu(x)
|
||||
rms = torch.ops.vllm_ir.rms_norm(x, self.weight, EPS)
|
||||
q1, s1 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(rms, GROUP_SIZE)
|
||||
q2, s2 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(rms, GROUP_SIZE)
|
||||
return q1, s1, q2, s2
|
||||
|
||||
|
||||
class _ViewDoubleQuantModel(torch.nn.Module):
|
||||
"""``rms_norm -> view -> 2x group_fp8_quant`` fan-out (DSv3.2 shape).
|
||||
|
||||
Reproduces the FX-graph shape produced by ``Fp8BlockScaledMMLinearKernel``'s
|
||||
2D-flatten before the FP8 group quant op.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.weight = torch.nn.Parameter(torch.ones(HIDDEN_SIZE, dtype=torch.bfloat16))
|
||||
|
||||
def forward(
|
||||
self, x: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
x = torch.relu(x)
|
||||
rms = torch.ops.vllm_ir.rms_norm(x, self.weight, EPS)
|
||||
view = rms.view(-1, rms.shape[-1])
|
||||
q1, s1 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(view, GROUP_SIZE)
|
||||
q2, s2 = torch.ops.vllm.rocm_aiter_group_fp8_quant.default(view, GROUP_SIZE)
|
||||
return q1, s1, q2, s2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_cls",
|
||||
[_NoViewDoubleQuantModel, _ViewDoubleQuantModel],
|
||||
ids=["no_view", "with_view"],
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
not is_aiter_found_and_supported(),
|
||||
reason="Only test on ROCm with AITER installed and supported",
|
||||
)
|
||||
def test_double_aiter_rms_fp8_group_quant_fusion(
|
||||
model_cls: type[torch.nn.Module],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""
|
||||
Both fan-out shapes (with and without an intermediate view) must fuse
|
||||
into ``rocm_aiter_rmsnorm_fp8_group_quant``: the no-view shape via
|
||||
``DoubleAiterRMSFp8GroupQuantPattern`` and the viewed shape via the
|
||||
new ``DoubleAiterRMSFp8GroupQuantViewPattern`` sibling.
|
||||
|
||||
A failure on the ``with_view`` parametrization is a regression on the
|
||||
DSv3.2 q_c norm path that this PR's view-tolerant pattern is intended
|
||||
to cover.
|
||||
"""
|
||||
torch._dynamo.reset()
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
model_config=ModelConfig(dtype=torch.bfloat16),
|
||||
compilation_config=CompilationConfig(
|
||||
mode=CompilationMode.VLLM_COMPILE,
|
||||
custom_ops=["+rms_norm", "+quant_fp8"],
|
||||
pass_config=PassConfig(
|
||||
fuse_norm_quant=True,
|
||||
eliminate_noops=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
|
||||
from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
|
||||
RocmAiterRMSNormQuantFusionPass,
|
||||
)
|
||||
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
torch.manual_seed(0)
|
||||
|
||||
m.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
|
||||
fusion_pass = RocmAiterRMSNormQuantFusionPass(vllm_config)
|
||||
passes = [
|
||||
NoOpEliminationPass(vllm_config),
|
||||
fusion_pass,
|
||||
PostCleanupPass(vllm_config),
|
||||
]
|
||||
backend = TestBackend(*passes)
|
||||
model = model_cls()
|
||||
|
||||
x = torch.randn(8, HIDDEN_SIZE)
|
||||
torch._dynamo.mark_dynamic(x, 0)
|
||||
|
||||
outputs_unfused = model(x)
|
||||
model_fused = torch.compile(model, backend=backend)
|
||||
outputs_fused = model_fused(x)
|
||||
|
||||
# Both consumers must be rewritten into the fused op (one
|
||||
# ``register_replacement`` rewrite covers the whole 1-to-2 fan-out).
|
||||
assert fusion_pass.matched_count == 1, (
|
||||
f"Expected the {model_cls.__name__} fan-out to fuse via the "
|
||||
f"DoubleQuant pattern (matched_count == 1), got "
|
||||
f"{fusion_pass.matched_count}"
|
||||
)
|
||||
|
||||
fused_op = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op()
|
||||
backend.check_after_ops([fused_op])
|
||||
|
||||
# Numerical parity sanity-check: the fused pair must match the
|
||||
# unfused pair on FP8 outputs (exact byte-equality is the goal,
|
||||
# but allow a tiny tolerance for any residual numeric noise).
|
||||
for fused_t, unfused_t in zip(outputs_fused, outputs_unfused):
|
||||
torch.testing.assert_close(
|
||||
fused_t.to(torch.float32),
|
||||
unfused_t.to(torch.float32),
|
||||
atol=1e-2,
|
||||
rtol=1e-2,
|
||||
)
|
||||
@@ -0,0 +1,413 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm.config
|
||||
from tests.compile.backend import TestBackend
|
||||
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
|
||||
from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
|
||||
from vllm.compilation.passes.fusion.mla_rope_kvcache_cat_fusion import (
|
||||
MLARoPEKVCacheCatFusionPass,
|
||||
)
|
||||
from vllm.compilation.passes.utility.fix_functionalization import (
|
||||
FixFunctionalizationPass,
|
||||
)
|
||||
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
|
||||
from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
|
||||
from vllm.config import (
|
||||
CacheConfig,
|
||||
CompilationConfig,
|
||||
CompilationMode,
|
||||
ModelConfig,
|
||||
PassConfig,
|
||||
VllmConfig,
|
||||
)
|
||||
from vllm.forward_context import get_forward_context, set_forward_context
|
||||
from vllm.model_executor.layers.attention import MLAAttention
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear
|
||||
from vllm.model_executor.layers.rotary_embedding import (
|
||||
DeepseekScalingRotaryEmbedding,
|
||||
RotaryEmbedding,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import _encode_layer_name
|
||||
from vllm.v1.attention.backend import (
|
||||
AttentionBackend,
|
||||
CommonAttentionMetadata,
|
||||
)
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
INDEX_SELECT_OP = torch.ops.aten.index.Tensor
|
||||
VLLM_UNIFIED_MLA_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_mla_kv_cache_update
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
|
||||
class MLARoPEKVCacheCatTestModel(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
attn_backend: AttentionBackendEnum,
|
||||
use_deepseek_scaling_rope: bool,
|
||||
num_heads: int,
|
||||
qk_nope_head_dim: int,
|
||||
qk_rope_head_dim: int,
|
||||
v_head_dim: int,
|
||||
q_lora_rank: int,
|
||||
kv_lora_rank: int,
|
||||
is_neox: bool,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
prefix: str = "model.layers.0.self_attn.attn",
|
||||
):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.qk_nope_head_dim = qk_nope_head_dim
|
||||
self.qk_rope_head_dim = qk_rope_head_dim
|
||||
self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
|
||||
self.v_head_dim = v_head_dim
|
||||
self.q_lora_rank = q_lora_rank
|
||||
self.kv_lora_rank = kv_lora_rank
|
||||
self.dtype = dtype
|
||||
self.device = device
|
||||
self.layer_name = prefix
|
||||
|
||||
self.num_kv_heads = 1
|
||||
self.head_size = kv_lora_rank + qk_rope_head_dim
|
||||
self.block_size = vllm_config.cache_config.block_size
|
||||
self.scale = self.qk_head_dim**-0.5
|
||||
|
||||
if use_deepseek_scaling_rope:
|
||||
self.rotary_emb = DeepseekScalingRotaryEmbedding(
|
||||
head_size=qk_rope_head_dim,
|
||||
rotary_dim=qk_rope_head_dim,
|
||||
max_position_embeddings=4096,
|
||||
base=10000,
|
||||
is_neox_style=is_neox,
|
||||
scaling_factor=1.0,
|
||||
dtype=dtype,
|
||||
)
|
||||
else:
|
||||
self.rotary_emb = RotaryEmbedding(
|
||||
head_size=qk_rope_head_dim,
|
||||
rotary_dim=qk_rope_head_dim,
|
||||
max_position_embeddings=4096,
|
||||
base=10000,
|
||||
is_neox_style=is_neox,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Initialize intermediate mm layers for unit test
|
||||
self.q_b_proj = ColumnParallelLinear(
|
||||
self.q_lora_rank,
|
||||
self.num_heads * self.qk_head_dim,
|
||||
bias=False,
|
||||
prefix=f"{prefix}.q_b_proj",
|
||||
).to(device)
|
||||
self.kv_b_proj = ColumnParallelLinear(
|
||||
self.kv_lora_rank,
|
||||
self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
|
||||
bias=False,
|
||||
prefix=f"{prefix}.kv_b_proj",
|
||||
).to(device)
|
||||
|
||||
# ColumnParallelLinear default init in bf16 with seed 0 produces
|
||||
# near-zero weights (7/4.7M nonzero), making the GEMM output almost
|
||||
# entirely zero and masking correctness bugs. Reinitialize to get
|
||||
# dense outputs.
|
||||
with torch.no_grad():
|
||||
torch.nn.init.normal_(self.q_b_proj.weight, std=0.02)
|
||||
torch.nn.init.normal_(self.kv_b_proj.weight, std=0.02)
|
||||
|
||||
# Register layer metadata for the fusion pass via MLAAttention
|
||||
self.mla_attn = MLAAttention(
|
||||
num_heads=self.num_heads,
|
||||
scale=self.scale,
|
||||
qk_nope_head_dim=self.qk_nope_head_dim,
|
||||
qk_rope_head_dim=self.qk_rope_head_dim,
|
||||
v_head_dim=self.v_head_dim,
|
||||
q_lora_rank=self.q_lora_rank,
|
||||
kv_lora_rank=self.kv_lora_rank,
|
||||
kv_b_proj=self.kv_b_proj,
|
||||
cache_config=vllm_config.cache_config,
|
||||
quant_config=vllm_config.quant_config,
|
||||
prefix=prefix,
|
||||
attn_backend=attn_backend.get_class(),
|
||||
)
|
||||
self.attn_backend: type[AttentionBackend] = self.mla_attn.get_attn_backend()
|
||||
self.mla_attn._k_scale = self.mla_attn._k_scale.to(device)
|
||||
self.mla_attn._v_scale = self.mla_attn._v_scale.to(device)
|
||||
|
||||
# Keep both the string dtype (for ops) and torch dtype (for tensors)
|
||||
self.kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
|
||||
self.kv_cache_dtype = (
|
||||
FP8_DTYPE if self.kv_cache_dtype_str.startswith("fp8") else self.dtype
|
||||
)
|
||||
|
||||
# Initialize attn MetadataBuilder
|
||||
self.builder = self.attn_backend.get_builder_cls()(
|
||||
kv_cache_spec=self.mla_attn.get_kv_cache_spec(vllm_config),
|
||||
layer_names=[self.mla_attn.layer_name],
|
||||
vllm_config=vllm_config,
|
||||
device=device,
|
||||
)
|
||||
|
||||
def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata:
|
||||
"""Initialize attention metadata."""
|
||||
# Create common attn metadata
|
||||
batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
|
||||
common_attn_metadata = create_common_attn_metadata(
|
||||
batch_spec, self.block_size, self.device, arange_block_indices=True
|
||||
)
|
||||
|
||||
max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
|
||||
num_blocks = batch_size * max_blocks
|
||||
|
||||
# Fetch the attention backend and kv cache shape and stride order
|
||||
kv_cache_shape = self.attn_backend.get_kv_cache_shape(
|
||||
num_blocks, self.block_size, self.num_kv_heads, self.head_size
|
||||
)
|
||||
try:
|
||||
kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order()
|
||||
except (AttributeError, NotImplementedError):
|
||||
kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
|
||||
|
||||
kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
|
||||
inv_order = [
|
||||
kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
|
||||
]
|
||||
|
||||
raw_tensor = torch.zeros(
|
||||
num_blocks * self.block_size * self.num_kv_heads * self.head_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
device=self.device,
|
||||
)
|
||||
raw_tensor = raw_tensor.view(kv_cache_shape)
|
||||
kv_cache = raw_tensor.permute(*inv_order)
|
||||
|
||||
self.mla_attn.kv_cache = kv_cache
|
||||
|
||||
# Build attn metadata
|
||||
attn_metadata = self.builder.build(
|
||||
common_prefix_len=0, common_attn_metadata=common_attn_metadata
|
||||
)
|
||||
|
||||
return attn_metadata
|
||||
|
||||
def forward(
|
||||
self, qkv_lora: torch.Tensor, positions: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
qkv_lora = qkv_lora.clone()
|
||||
q_c, kv_lora = qkv_lora.split(
|
||||
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
|
||||
dim=-1,
|
||||
)
|
||||
q = self.q_b_proj(q_c)[0]
|
||||
kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
|
||||
|
||||
q = q.view(-1, self.num_heads, self.qk_head_dim)
|
||||
k_pe = k_pe.unsqueeze(1)
|
||||
|
||||
q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb(
|
||||
positions, q[..., self.qk_nope_head_dim :], k_pe
|
||||
)
|
||||
|
||||
dummy = torch.ops.vllm.unified_mla_kv_cache_update(
|
||||
kv_c,
|
||||
k_pe,
|
||||
_encode_layer_name(self.layer_name),
|
||||
self.kv_cache_dtype_str,
|
||||
self.mla_attn._k_scale,
|
||||
)
|
||||
return q, kv_c, k_pe, dummy
|
||||
|
||||
def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
|
||||
ops = [
|
||||
INDEX_SELECT_OP,
|
||||
torch.ops.vllm.unified_mla_kv_cache_update.default,
|
||||
]
|
||||
return ops
|
||||
|
||||
def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
|
||||
return [torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default]
|
||||
|
||||
|
||||
MLA_BACKENDS = [AttentionBackendEnum.TRITON_MLA]
|
||||
if flash_attn_supports_mla():
|
||||
MLA_BACKENDS += [AttentionBackendEnum.FLASH_ATTN_MLA]
|
||||
if is_aiter_found_and_supported():
|
||||
MLA_BACKENDS += [AttentionBackendEnum.ROCM_AITER_MLA]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("attn_backend", MLA_BACKENDS)
|
||||
@pytest.mark.parametrize("use_deepseek_scaling_rope", [True])
|
||||
@pytest.mark.parametrize("num_heads", [16])
|
||||
@pytest.mark.parametrize("qk_nope_head_dim", [128])
|
||||
@pytest.mark.parametrize("qk_rope_head_dim", [64])
|
||||
@pytest.mark.parametrize("v_head_dim", [128])
|
||||
@pytest.mark.parametrize("q_lora_rank", [1536])
|
||||
@pytest.mark.parametrize("kv_lora_rank", [512])
|
||||
@pytest.mark.parametrize("block_size", [16])
|
||||
@pytest.mark.parametrize("is_neox", [True, False])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="MLA RoPE+KVCache+Cat fusion is only supported on CUDA and ROCm.",
|
||||
)
|
||||
def test_mla_rope_kvcache_cat_fusion(
|
||||
attn_backend: AttentionBackendEnum,
|
||||
use_deepseek_scaling_rope: bool,
|
||||
num_heads: int,
|
||||
qk_nope_head_dim: int,
|
||||
qk_rope_head_dim: int,
|
||||
v_head_dim: int,
|
||||
q_lora_rank: int,
|
||||
kv_lora_rank: int,
|
||||
block_size: int,
|
||||
is_neox: bool,
|
||||
dtype: torch.dtype,
|
||||
kv_cache_dtype: str,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
torch.set_default_device("cuda")
|
||||
torch.set_default_dtype(dtype)
|
||||
torch.manual_seed(0)
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
model_config=ModelConfig(
|
||||
model="deepseek-ai/DeepSeek-V2-Lite",
|
||||
dtype=dtype,
|
||||
),
|
||||
cache_config=CacheConfig(
|
||||
block_size=block_size,
|
||||
cache_dtype=kv_cache_dtype,
|
||||
),
|
||||
compilation_config=CompilationConfig(
|
||||
mode=CompilationMode.VLLM_COMPILE,
|
||||
pass_config=PassConfig(
|
||||
fuse_rope_kvcache_cat_mla=True,
|
||||
eliminate_noops=True,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
|
||||
if not torch.distributed.is_initialized():
|
||||
from vllm.distributed.parallel_state import (
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": "0",
|
||||
"LOCAL_RANK": "0",
|
||||
"WORLD_SIZE": "1",
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": "54321",
|
||||
}
|
||||
)
|
||||
init_distributed_environment()
|
||||
initialize_model_parallel()
|
||||
|
||||
if attn_backend == AttentionBackendEnum.ROCM_AITER_MLA:
|
||||
m.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
|
||||
model = MLARoPEKVCacheCatTestModel(
|
||||
vllm_config=vllm_config,
|
||||
attn_backend=attn_backend,
|
||||
use_deepseek_scaling_rope=use_deepseek_scaling_rope,
|
||||
num_heads=num_heads,
|
||||
qk_nope_head_dim=qk_nope_head_dim,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
v_head_dim=v_head_dim,
|
||||
q_lora_rank=q_lora_rank,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
is_neox=is_neox,
|
||||
dtype=dtype,
|
||||
device=torch.get_default_device(),
|
||||
)
|
||||
|
||||
fusion_pass = MLARoPEKVCacheCatFusionPass(vllm_config)
|
||||
# note: FixFunctionalizationPass is required to correctly lower
|
||||
# the fused op to its inplace version with auto-functionalization v1.
|
||||
# Without it, decompose_auto_functionalized calls clone_preserve_strides
|
||||
# on the non-contiguous q_pe slice directly, and inductor's lowering
|
||||
# of the resulting as_strided chain incorrectly drops the storage offset.
|
||||
# auto-functionalization v2 avoids this: it clones the contiguous base
|
||||
# tensor (_all_bases) and reconstructs the slice as a view, so the
|
||||
# offset is never passed through as_strided lowering.
|
||||
passes = [
|
||||
NoOpEliminationPass(vllm_config),
|
||||
fusion_pass,
|
||||
PostCleanupPass(vllm_config),
|
||||
FixFunctionalizationPass(vllm_config),
|
||||
]
|
||||
backend = TestBackend(*passes)
|
||||
|
||||
T = 5
|
||||
|
||||
qkv_lora = torch.randn(
|
||||
T,
|
||||
q_lora_rank + kv_lora_rank + qk_rope_head_dim,
|
||||
dtype=dtype,
|
||||
)
|
||||
pos = torch.arange(T, dtype=torch.long)
|
||||
|
||||
qkv_unfused = qkv_lora.clone()
|
||||
pos_unfused = pos.clone()
|
||||
|
||||
# Run unfused version
|
||||
with set_forward_context(None, vllm_config):
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata = model.build_attn_metadata(T)
|
||||
forward_context.slot_mapping = {
|
||||
model.layer_name: attn_metadata.slot_mapping
|
||||
}
|
||||
q_unfused, kv_c_unfused, k_pe_unfused, dummy = model(
|
||||
qkv_unfused, pos_unfused
|
||||
)
|
||||
attn_layer = forward_context.no_compile_layers[model.layer_name]
|
||||
kv_cache_unfused = attn_layer.kv_cache.clone()
|
||||
del dummy
|
||||
|
||||
# Run fused version (compiled)
|
||||
torch._dynamo.mark_dynamic(qkv_lora, 0)
|
||||
torch._dynamo.mark_dynamic(pos, 0)
|
||||
with set_forward_context(None, vllm_config):
|
||||
model_fused = torch.compile(model, backend=backend)
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata = model.build_attn_metadata(T)
|
||||
forward_context.slot_mapping = {
|
||||
model.layer_name: attn_metadata.slot_mapping
|
||||
}
|
||||
q_fused, kv_c_fused, k_pe_fused, dummy = model_fused(qkv_lora, pos)
|
||||
attn_layer = forward_context.no_compile_layers[model.layer_name]
|
||||
kv_cache_fused = attn_layer.kv_cache
|
||||
del dummy
|
||||
|
||||
assert fusion_pass.matched_count == 1
|
||||
|
||||
backend.check_before_ops(model.ops_in_model_before())
|
||||
backend.check_after_ops(model.ops_in_model_after())
|
||||
|
||||
if dtype == torch.float16:
|
||||
ATOL, RTOL = (2e-3, 2e-3)
|
||||
else:
|
||||
ATOL, RTOL = (1e-2, 1e-2)
|
||||
|
||||
torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
|
||||
torch.testing.assert_close(kv_c_unfused, kv_c_fused, atol=ATOL, rtol=RTOL)
|
||||
torch.testing.assert_close(k_pe_unfused, k_pe_fused, atol=ATOL, rtol=RTOL)
|
||||
# Cannot compare fp8_* directly here, cast to model dtype instead
|
||||
torch.testing.assert_close(
|
||||
kv_cache_unfused.view(dtype),
|
||||
kv_cache_fused.view(dtype),
|
||||
atol=ATOL,
|
||||
rtol=RTOL,
|
||||
)
|
||||
@@ -34,7 +34,6 @@ from vllm.v1.attention.backend import (
|
||||
CommonAttentionMetadata,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.kv_cache_interface import AttentionSpec
|
||||
|
||||
INDEX_SELECT_OP = torch.ops.aten.index.Tensor
|
||||
VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
|
||||
@@ -102,13 +101,8 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
|
||||
)
|
||||
|
||||
# Initialize attn MetadataBuilder
|
||||
self.builder = self.attn.attn_backend.get_builder_cls()(
|
||||
kv_cache_spec=AttentionSpec(
|
||||
block_size=self.block_size,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
head_size=head_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
),
|
||||
self.builder = self.attn_backend.get_builder_cls()(
|
||||
kv_cache_spec=self.attn.get_kv_cache_spec(vllm_config),
|
||||
layer_names=[self.attn.layer_name],
|
||||
vllm_config=vllm_config,
|
||||
device=device,
|
||||
@@ -126,12 +120,11 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
|
||||
num_blocks = batch_size * max_blocks
|
||||
|
||||
# Fetch the attention backend and kv cache shape and stride order
|
||||
attn_backend = self.attn.attn_backend
|
||||
kv_cache_shape = attn_backend.get_kv_cache_shape(
|
||||
kv_cache_shape = self.attn_backend.get_kv_cache_shape(
|
||||
num_blocks, self.block_size, self.num_kv_heads, self.head_size
|
||||
)
|
||||
try:
|
||||
kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
|
||||
kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order()
|
||||
except (AttributeError, NotImplementedError):
|
||||
kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
|
||||
|
||||
|
||||
+12
-21
@@ -224,28 +224,19 @@ def init_test_http_connection():
|
||||
def dist_init():
|
||||
from tests.utils import ensure_current_vllm_config
|
||||
|
||||
# Close the fd returned by mkstemp; FileStore opens the path itself.
|
||||
# Leaving it open leaks one FD per test and eventually exhausts the
|
||||
# ulimit, causing FileStore's destructor to throw c10::DistStoreError
|
||||
# ("Too many open files") during gc and abort the process.
|
||||
fd, temp_file = tempfile.mkstemp()
|
||||
os.close(fd)
|
||||
temp_file = tempfile.mkstemp()[1]
|
||||
|
||||
try:
|
||||
with ensure_current_vllm_config():
|
||||
init_distributed_environment(
|
||||
world_size=1,
|
||||
rank=0,
|
||||
distributed_init_method=f"file://{temp_file}",
|
||||
local_rank=0,
|
||||
backend="nccl",
|
||||
)
|
||||
initialize_model_parallel(1, 1)
|
||||
yield
|
||||
cleanup_dist_env_and_memory()
|
||||
finally:
|
||||
with contextlib.suppress(OSError):
|
||||
os.unlink(temp_file)
|
||||
with ensure_current_vllm_config():
|
||||
init_distributed_environment(
|
||||
world_size=1,
|
||||
rank=0,
|
||||
distributed_init_method=f"file://{temp_file}",
|
||||
local_rank=0,
|
||||
backend="nccl",
|
||||
)
|
||||
initialize_model_parallel(1, 1)
|
||||
yield
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
||||
@@ -67,58 +67,3 @@ async def test_chat_with_enable_force_include_usage(
|
||||
chunk.usage.prompt_tokens + chunk.usage.completion_tokens
|
||||
)
|
||||
last_completion_tokens = chunk.usage.completion_tokens
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def transcription_server_with_force_include_usage():
|
||||
args = [
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-num-seqs",
|
||||
"4",
|
||||
"--enforce-eager",
|
||||
"--enable-force-include-usage",
|
||||
"--gpu-memory-utilization",
|
||||
"0.2",
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer("openai/whisper-large-v3-turbo", args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def transcription_client_with_force_include_usage(
|
||||
transcription_server_with_force_include_usage,
|
||||
):
|
||||
async with (
|
||||
transcription_server_with_force_include_usage.get_async_client() as async_client
|
||||
):
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_transcription_with_enable_force_include_usage(
|
||||
transcription_client_with_force_include_usage, winning_call
|
||||
):
|
||||
res = (
|
||||
await transcription_client_with_force_include_usage.audio.transcriptions.create(
|
||||
model="openai/whisper-large-v3-turbo",
|
||||
file=winning_call,
|
||||
language="en",
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
timeout=30,
|
||||
)
|
||||
)
|
||||
|
||||
async for chunk in res:
|
||||
if not len(chunk.choices):
|
||||
# final usage sent
|
||||
usage = chunk.usage
|
||||
assert isinstance(usage, dict)
|
||||
assert usage["prompt_tokens"] > 0
|
||||
assert usage["completion_tokens"] > 0
|
||||
assert usage["total_tokens"] > 0
|
||||
else:
|
||||
assert not hasattr(chunk, "usage")
|
||||
|
||||
@@ -0,0 +1,492 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Cross-API render parity tests.
|
||||
|
||||
Verifies that the chat completion input path (parse_chat_input_to_harmony_message)
|
||||
and the responses API input path (response_input_to_harmony) produce identical
|
||||
Harmony messages and identical rendered token sequences when given equivalent
|
||||
conversation representations.
|
||||
|
||||
The chat completion API encodes reasoning and tool calls as fields on a single
|
||||
assistant message dict; the responses API encodes them as separate typed items
|
||||
in request.input. Both paths must converge on the same Harmony message list and
|
||||
therefore the same rendered prompt.
|
||||
|
||||
Each test:
|
||||
1. Builds Harmony messages from each path for a single message or sequence.
|
||||
2. Asserts message-level properties (role, channel, recipient, content)
|
||||
using verify_harmony_messages.
|
||||
3. Asserts that render_for_completion produces identical token sequences.
|
||||
"""
|
||||
|
||||
from openai.types.responses import ResponseFunctionToolCall
|
||||
|
||||
from tests.entrypoints.openai.utils import verify_harmony_messages
|
||||
from vllm.entrypoints.openai.parser.harmony_utils import (
|
||||
get_system_message,
|
||||
parse_chat_input_to_harmony_message,
|
||||
render_for_completion,
|
||||
)
|
||||
from vllm.entrypoints.openai.responses.harmony import response_input_to_harmony
|
||||
|
||||
# Use a fixed date so the system message is deterministic across both paths.
|
||||
_DATE = "2025-01-01"
|
||||
|
||||
|
||||
def _system():
|
||||
return get_system_message(start_date=_DATE)
|
||||
|
||||
|
||||
class TestResponseInputToHarmonyRenderParity:
|
||||
"""Each test drives the same conversation through both APIs and asserts
|
||||
identical Harmony messages and rendered token sequences."""
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Single-message cases
|
||||
# -----------------------------------------------------------------------
|
||||
|
||||
def test_user_message(self):
|
||||
chat_msgs = parse_chat_input_to_harmony_message(
|
||||
{"role": "user", "content": "What's the weather in Paris?"}
|
||||
)
|
||||
resp_msgs = [
|
||||
response_input_to_harmony(
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": "What's the weather in Paris?",
|
||||
},
|
||||
prev_responses=[],
|
||||
)
|
||||
]
|
||||
|
||||
expected = [{"role": "user", "content": "What's the weather in Paris?"}]
|
||||
verify_harmony_messages(chat_msgs, expected)
|
||||
verify_harmony_messages(resp_msgs, expected)
|
||||
|
||||
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
|
||||
[_system()] + resp_msgs
|
||||
)
|
||||
|
||||
def test_assistant_final_message(self):
|
||||
chat_msgs = parse_chat_input_to_harmony_message(
|
||||
{"role": "assistant", "content": "It is 18°C in Paris."}
|
||||
)
|
||||
resp_msgs = [
|
||||
response_input_to_harmony(
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": "It is 18°C in Paris.",
|
||||
},
|
||||
prev_responses=[],
|
||||
)
|
||||
]
|
||||
|
||||
expected = [
|
||||
{"role": "assistant", "channel": "final", "content": "It is 18°C in Paris."}
|
||||
]
|
||||
verify_harmony_messages(chat_msgs, expected)
|
||||
verify_harmony_messages(resp_msgs, expected)
|
||||
|
||||
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
|
||||
[_system()] + resp_msgs
|
||||
)
|
||||
|
||||
def test_reasoning_item(self):
|
||||
# Chat path: assistant message with only a reasoning field and no content.
|
||||
chat_msgs = parse_chat_input_to_harmony_message(
|
||||
{
|
||||
"role": "assistant",
|
||||
"reasoning": "I should call get_weather.",
|
||||
"content": "",
|
||||
}
|
||||
)
|
||||
resp_msgs = [
|
||||
response_input_to_harmony(
|
||||
{
|
||||
"type": "reasoning",
|
||||
"content": [
|
||||
{"type": "reasoning_text", "text": "I should call get_weather."}
|
||||
],
|
||||
},
|
||||
prev_responses=[],
|
||||
)
|
||||
]
|
||||
|
||||
expected = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"channel": "analysis",
|
||||
"content": "I should call get_weather.",
|
||||
}
|
||||
]
|
||||
verify_harmony_messages(chat_msgs, expected)
|
||||
verify_harmony_messages(resp_msgs, expected)
|
||||
|
||||
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
|
||||
[_system()] + resp_msgs
|
||||
)
|
||||
|
||||
def test_function_call(self):
|
||||
chat_msgs = parse_chat_input_to_harmony_message(
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"location": "Paris"}',
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
resp_msgs = [
|
||||
response_input_to_harmony(
|
||||
{
|
||||
"type": "function_call",
|
||||
"name": "get_weather",
|
||||
"arguments": '{"location": "Paris"}',
|
||||
},
|
||||
prev_responses=[],
|
||||
)
|
||||
]
|
||||
|
||||
expected = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"channel": "commentary",
|
||||
"recipient": "functions.get_weather",
|
||||
"content": '{"location": "Paris"}',
|
||||
"content_type": "json",
|
||||
}
|
||||
]
|
||||
verify_harmony_messages(chat_msgs, expected)
|
||||
verify_harmony_messages(resp_msgs, expected)
|
||||
|
||||
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
|
||||
[_system()] + resp_msgs
|
||||
)
|
||||
|
||||
def test_tool_output(self):
|
||||
prev_call = ResponseFunctionToolCall(
|
||||
id="fc_1",
|
||||
call_id="call_1",
|
||||
name="get_weather",
|
||||
arguments='{"location": "Paris"}',
|
||||
type="function_call",
|
||||
)
|
||||
|
||||
chat_msgs = parse_chat_input_to_harmony_message(
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "18°C, clear skies."},
|
||||
tool_id_names={"call_1": "get_weather"},
|
||||
)
|
||||
resp_msgs = [
|
||||
response_input_to_harmony(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_1",
|
||||
"output": "18°C, clear skies.",
|
||||
},
|
||||
prev_responses=[prev_call],
|
||||
)
|
||||
]
|
||||
|
||||
expected = [
|
||||
{
|
||||
"role": "tool",
|
||||
"author_name": "functions.get_weather",
|
||||
"channel": "commentary",
|
||||
"recipient": "assistant",
|
||||
"content": "18°C, clear skies.",
|
||||
}
|
||||
]
|
||||
verify_harmony_messages(chat_msgs, expected)
|
||||
verify_harmony_messages(resp_msgs, expected)
|
||||
|
||||
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
|
||||
[_system()] + resp_msgs
|
||||
)
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Combined and multi-turn cases
|
||||
# -----------------------------------------------------------------------
|
||||
|
||||
def test_reasoning_combined_with_function_call(self):
|
||||
"""Chat API packs reasoning + tool_calls into one dict; responses API
|
||||
represents them as two separate items. Both must produce the same two
|
||||
Harmony messages in the same order: analysis then commentary."""
|
||||
chat_msgs = parse_chat_input_to_harmony_message(
|
||||
{
|
||||
"role": "assistant",
|
||||
"reasoning": "I should get the weather for Paris.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"location": "Paris"}',
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
resp_msgs = [
|
||||
response_input_to_harmony(
|
||||
{
|
||||
"type": "reasoning",
|
||||
"content": [
|
||||
{
|
||||
"type": "reasoning_text",
|
||||
"text": "I should get the weather for Paris.",
|
||||
}
|
||||
],
|
||||
},
|
||||
prev_responses=[],
|
||||
),
|
||||
response_input_to_harmony(
|
||||
{
|
||||
"type": "function_call",
|
||||
"name": "get_weather",
|
||||
"arguments": '{"location": "Paris"}',
|
||||
},
|
||||
prev_responses=[],
|
||||
),
|
||||
]
|
||||
|
||||
expected = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"channel": "analysis",
|
||||
"content": "I should get the weather for Paris.",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"channel": "commentary",
|
||||
"recipient": "functions.get_weather",
|
||||
"content": '{"location": "Paris"}',
|
||||
"content_type": "json",
|
||||
},
|
||||
]
|
||||
verify_harmony_messages(chat_msgs, expected)
|
||||
verify_harmony_messages(resp_msgs, expected)
|
||||
|
||||
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
|
||||
[_system()] + resp_msgs
|
||||
)
|
||||
|
||||
def test_full_multi_turn_tool_call_conversation(self):
|
||||
"""Full conversation: user -> reasoning + tool_call -> tool_output -> final.
|
||||
|
||||
Both APIs must render the complete conversation to identical token sequences.
|
||||
This exercises the entire input pipeline including all message types and
|
||||
the Rust harmony encoder.
|
||||
"""
|
||||
prev_call = ResponseFunctionToolCall(
|
||||
id="fc_1",
|
||||
call_id="call_1",
|
||||
name="get_weather",
|
||||
arguments='{"location": "Paris"}',
|
||||
type="function_call",
|
||||
)
|
||||
|
||||
# --- Chat completion API path ---
|
||||
tool_id_names = {"call_1": "get_weather"}
|
||||
chat_msgs = []
|
||||
chat_msgs += parse_chat_input_to_harmony_message(
|
||||
{"role": "user", "content": "What's the weather in Paris?"}
|
||||
)
|
||||
chat_msgs += parse_chat_input_to_harmony_message(
|
||||
{
|
||||
"role": "assistant",
|
||||
"reasoning": "I should call get_weather for Paris.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"location": "Paris"}',
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
chat_msgs += parse_chat_input_to_harmony_message(
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "18°C, clear skies."},
|
||||
tool_id_names=tool_id_names,
|
||||
)
|
||||
chat_msgs += parse_chat_input_to_harmony_message(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "It is currently 18°C in Paris with clear skies.",
|
||||
}
|
||||
)
|
||||
|
||||
# --- Responses API path ---
|
||||
resp_input = [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": "What's the weather in Paris?",
|
||||
},
|
||||
{
|
||||
"type": "reasoning",
|
||||
"content": [
|
||||
{
|
||||
"type": "reasoning_text",
|
||||
"text": "I should call get_weather for Paris.",
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"type": "function_call",
|
||||
"name": "get_weather",
|
||||
"arguments": '{"location": "Paris"}',
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_1",
|
||||
"output": "18°C, clear skies.",
|
||||
},
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": "It is currently 18°C in Paris with clear skies.",
|
||||
},
|
||||
]
|
||||
resp_msgs = [
|
||||
response_input_to_harmony(item, prev_responses=[prev_call])
|
||||
for item in resp_input
|
||||
]
|
||||
|
||||
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
|
||||
[_system()] + resp_msgs
|
||||
)
|
||||
|
||||
def test_multi_turn_two_tool_calls_with_reasoning_between(self):
|
||||
"""Validates parity for a chain of two tool calls, each with its own
|
||||
reasoning trace. Reasoning traces in between commentary-channel tool
|
||||
calls must survive as analysis-channel messages in both paths.
|
||||
"""
|
||||
prev_call_1 = ResponseFunctionToolCall(
|
||||
id="fc_1",
|
||||
call_id="call_1",
|
||||
name="get_weather",
|
||||
arguments='{"location": "Paris"}',
|
||||
type="function_call",
|
||||
)
|
||||
prev_call_2 = ResponseFunctionToolCall(
|
||||
id="fc_2",
|
||||
call_id="call_2",
|
||||
name="get_forecast",
|
||||
arguments='{"location": "Paris", "days": 7}',
|
||||
type="function_call",
|
||||
)
|
||||
|
||||
# --- Chat completion API path ---
|
||||
tool_id_names = {"call_1": "get_weather", "call_2": "get_forecast"}
|
||||
chat_msgs = []
|
||||
chat_msgs += parse_chat_input_to_harmony_message(
|
||||
{"role": "user", "content": "What's the weather and forecast for Paris?"}
|
||||
)
|
||||
# First reasoning + tool call
|
||||
chat_msgs += parse_chat_input_to_harmony_message(
|
||||
{
|
||||
"role": "assistant",
|
||||
"reasoning": "I need current weather first.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"location": "Paris"}',
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
chat_msgs += parse_chat_input_to_harmony_message(
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "18°C, clear skies."},
|
||||
tool_id_names=tool_id_names,
|
||||
)
|
||||
# Second reasoning + tool call
|
||||
chat_msgs += parse_chat_input_to_harmony_message(
|
||||
{
|
||||
"role": "assistant",
|
||||
"reasoning": "Now I need the weekly forecast.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_2",
|
||||
"function": {
|
||||
"name": "get_forecast",
|
||||
"arguments": '{"location": "Paris", "days": 7}',
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
chat_msgs += parse_chat_input_to_harmony_message(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_2",
|
||||
"content": "Mon 17°C, Tue 19°C, Wed 16°C",
|
||||
},
|
||||
tool_id_names=tool_id_names,
|
||||
)
|
||||
|
||||
# --- Responses API path ---
|
||||
prev_responses = [prev_call_1, prev_call_2]
|
||||
resp_input = [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": "What's the weather and forecast for Paris?",
|
||||
},
|
||||
# First reasoning + tool call
|
||||
{
|
||||
"type": "reasoning",
|
||||
"content": [
|
||||
{"type": "reasoning_text", "text": "I need current weather first."}
|
||||
],
|
||||
},
|
||||
{
|
||||
"type": "function_call",
|
||||
"name": "get_weather",
|
||||
"arguments": '{"location": "Paris"}',
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_1",
|
||||
"output": "18°C, clear skies.",
|
||||
},
|
||||
# Second reasoning + tool call
|
||||
{
|
||||
"type": "reasoning",
|
||||
"content": [
|
||||
{
|
||||
"type": "reasoning_text",
|
||||
"text": "Now I need the weekly forecast.",
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"type": "function_call",
|
||||
"name": "get_forecast",
|
||||
"arguments": '{"location": "Paris", "days": 7}',
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_2",
|
||||
"output": "Mon 17°C, Tue 19°C, Wed 16°C",
|
||||
},
|
||||
]
|
||||
resp_msgs = [
|
||||
response_input_to_harmony(item, prev_responses=prev_responses)
|
||||
for item in resp_input
|
||||
]
|
||||
|
||||
assert render_for_completion([_system()] + chat_msgs) == render_for_completion(
|
||||
[_system()] + resp_msgs
|
||||
)
|
||||
@@ -325,8 +325,12 @@ async def test_function_calling_with_streaming_expected_arguments(
|
||||
"tool_choice",
|
||||
["auto", "required", {"type": "function", "name": "get_current_weather"}],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"enable_thinking",
|
||||
[True, False],
|
||||
)
|
||||
async def test_function_calling_with_streaming_types(
|
||||
client: openai.AsyncOpenAI, model_name: str, tool_choice
|
||||
client: openai.AsyncOpenAI, model_name: str, tool_choice, enable_thinking: bool
|
||||
):
|
||||
# this links the "done" type with the "start" type
|
||||
# so every "done" type should have a corresponding "start" type
|
||||
@@ -436,6 +440,7 @@ async def test_function_calling_with_streaming_types(
|
||||
input=input_list,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": enable_thinking}},
|
||||
stream=True,
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,274 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for response_input_to_harmony.
|
||||
|
||||
Covers every type branch in the function and verifies that each produced
|
||||
Harmony Message has the correct role, channel, recipient, content_type,
|
||||
author name, and text content.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from openai.types.responses import ResponseFunctionToolCall, ResponseReasoningItem
|
||||
from openai.types.responses.response_reasoning_item import (
|
||||
Content as ReasoningTextContent,
|
||||
)
|
||||
from openai_harmony import Role
|
||||
|
||||
from vllm.entrypoints.openai.responses.harmony import response_input_to_harmony
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PREV_CALL = ResponseFunctionToolCall(
|
||||
id="fc_test",
|
||||
call_id="call_test",
|
||||
name="get_weather",
|
||||
arguments='{"location": "Paris"}',
|
||||
type="function_call",
|
||||
)
|
||||
|
||||
_REASONING_ITEM = ResponseReasoningItem(
|
||||
id="rs_test",
|
||||
type="reasoning",
|
||||
content=[ReasoningTextContent(type="reasoning_text", text="Thinking hard.")],
|
||||
summary=[],
|
||||
status=None,
|
||||
)
|
||||
|
||||
|
||||
class TestResponseInputToHarmonyMessage:
|
||||
"""Unit tests for every message type handled by response_input_to_harmony."""
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# type="message" (or no type key)
|
||||
# -----------------------------------------------------------------------
|
||||
|
||||
def test_user_message_string_content(self):
|
||||
msg = response_input_to_harmony(
|
||||
{"type": "message", "role": "user", "content": "Hello"},
|
||||
prev_responses=[],
|
||||
)
|
||||
|
||||
assert msg.author.role == Role.USER
|
||||
assert msg.content[0].text == "Hello"
|
||||
assert msg.channel is None
|
||||
|
||||
def test_no_type_key_defaults_to_message_branch(self):
|
||||
"""Omitting 'type' should fall through to the message branch."""
|
||||
msg = response_input_to_harmony(
|
||||
{"role": "user", "content": "Hello"},
|
||||
prev_responses=[],
|
||||
)
|
||||
|
||||
assert msg.author.role == Role.USER
|
||||
assert msg.content[0].text == "Hello"
|
||||
|
||||
def test_system_message(self):
|
||||
msg = response_input_to_harmony(
|
||||
{"type": "message", "role": "system", "content": "Be helpful."},
|
||||
prev_responses=[],
|
||||
)
|
||||
|
||||
assert msg.author.role == Role.SYSTEM
|
||||
assert msg.content[0].text == "Be helpful."
|
||||
assert msg.channel is None
|
||||
|
||||
def test_assistant_message_gets_final_channel(self):
|
||||
msg = response_input_to_harmony(
|
||||
{"type": "message", "role": "assistant", "content": "The answer is 42."},
|
||||
prev_responses=[],
|
||||
)
|
||||
|
||||
assert msg.author.role == Role.ASSISTANT
|
||||
assert msg.channel == "final"
|
||||
assert msg.content[0].text == "The answer is 42."
|
||||
|
||||
def test_developer_message_gets_instructions_prefix(self):
|
||||
msg = response_input_to_harmony(
|
||||
{"type": "message", "role": "developer", "content": "Be concise."},
|
||||
prev_responses=[],
|
||||
)
|
||||
|
||||
assert msg.author.role == Role.DEVELOPER
|
||||
assert msg.content[0].text == "Instructions:\nBe concise."
|
||||
assert msg.channel is None
|
||||
|
||||
def test_message_with_array_content(self):
|
||||
msg = response_input_to_harmony(
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Part one. "},
|
||||
{"type": "text", "text": "Part two."},
|
||||
],
|
||||
},
|
||||
prev_responses=[],
|
||||
)
|
||||
|
||||
assert msg.author.role == Role.USER
|
||||
assert len(msg.content) == 2
|
||||
assert msg.content[0].text == "Part one. "
|
||||
assert msg.content[1].text == "Part two."
|
||||
|
||||
def test_developer_message_array_content_gets_prefix_on_each_part(self):
|
||||
msg = response_input_to_harmony(
|
||||
{
|
||||
"type": "message",
|
||||
"role": "developer",
|
||||
"content": [
|
||||
{"type": "text", "text": "Rule 1."},
|
||||
{"type": "text", "text": "Rule 2."},
|
||||
],
|
||||
},
|
||||
prev_responses=[],
|
||||
)
|
||||
|
||||
assert msg.content[0].text == "Instructions:\nRule 1."
|
||||
assert msg.content[1].text == "Instructions:\nRule 2."
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# type="reasoning"
|
||||
# -----------------------------------------------------------------------
|
||||
|
||||
def test_reasoning_gets_analysis_channel(self):
|
||||
msg = response_input_to_harmony(
|
||||
{
|
||||
"type": "reasoning",
|
||||
"content": [
|
||||
{"type": "reasoning_text", "text": "I should call get_weather."}
|
||||
],
|
||||
},
|
||||
prev_responses=[],
|
||||
)
|
||||
|
||||
assert msg.author.role == Role.ASSISTANT
|
||||
assert msg.channel == "analysis"
|
||||
assert msg.content[0].text == "I should call get_weather."
|
||||
|
||||
def test_reasoning_pydantic_model_input(self):
|
||||
"""A Pydantic ResponseReasoningItem should be model_dump()'d before parsing."""
|
||||
msg = response_input_to_harmony(_REASONING_ITEM, prev_responses=[])
|
||||
|
||||
assert msg.author.role == Role.ASSISTANT
|
||||
assert msg.channel == "analysis"
|
||||
assert msg.content[0].text == "Thinking hard."
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# type="function_call"
|
||||
# -----------------------------------------------------------------------
|
||||
|
||||
def test_function_call_channel_recipient_and_content_type(self):
|
||||
msg = response_input_to_harmony(
|
||||
{
|
||||
"type": "function_call",
|
||||
"name": "get_weather",
|
||||
"arguments": '{"location": "Paris"}',
|
||||
},
|
||||
prev_responses=[],
|
||||
)
|
||||
|
||||
assert msg.author.role == Role.ASSISTANT
|
||||
assert msg.channel == "commentary"
|
||||
assert msg.recipient == "functions.get_weather"
|
||||
assert msg.content_type == "json"
|
||||
assert msg.content[0].text == '{"location": "Paris"}'
|
||||
|
||||
def test_function_call_empty_arguments(self):
|
||||
msg = response_input_to_harmony(
|
||||
{"type": "function_call", "name": "ping", "arguments": ""},
|
||||
prev_responses=[],
|
||||
)
|
||||
|
||||
assert msg.recipient == "functions.ping"
|
||||
assert msg.content[0].text == ""
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# type="function_call_output"
|
||||
# -----------------------------------------------------------------------
|
||||
|
||||
def test_function_call_output_channel_recipient_and_author_name(self):
|
||||
msg = response_input_to_harmony(
|
||||
{"type": "function_call_output", "call_id": "call_test", "output": "18°C"},
|
||||
prev_responses=[_PREV_CALL],
|
||||
)
|
||||
|
||||
assert msg.author.role == Role.TOOL
|
||||
assert msg.author.name == "functions.get_weather"
|
||||
assert msg.channel == "commentary"
|
||||
assert msg.recipient == "assistant"
|
||||
assert msg.content[0].text == "18°C"
|
||||
|
||||
def test_function_call_output_uses_most_recent_matching_call(self):
|
||||
"""When multiple prev_responses share a call_id, the last one wins
|
||||
because the search is reversed."""
|
||||
earlier = ResponseFunctionToolCall(
|
||||
id="fc_old",
|
||||
call_id="call_test",
|
||||
name="old_func",
|
||||
arguments="{}",
|
||||
type="function_call",
|
||||
)
|
||||
later = ResponseFunctionToolCall(
|
||||
id="fc_new",
|
||||
call_id="call_test",
|
||||
name="get_weather",
|
||||
arguments="{}",
|
||||
type="function_call",
|
||||
)
|
||||
|
||||
msg = response_input_to_harmony(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_test",
|
||||
"output": "result",
|
||||
},
|
||||
prev_responses=[earlier, later],
|
||||
)
|
||||
|
||||
assert msg.author.name == "functions.get_weather"
|
||||
|
||||
def test_function_call_output_skips_non_function_call_items_in_prev_responses(
|
||||
self,
|
||||
):
|
||||
"""ResponseReasoningItem entries in prev_responses should be ignored."""
|
||||
msg = response_input_to_harmony(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_test",
|
||||
"output": "18°C",
|
||||
},
|
||||
prev_responses=[_REASONING_ITEM, _PREV_CALL],
|
||||
)
|
||||
|
||||
assert msg.author.name == "functions.get_weather"
|
||||
|
||||
def test_function_call_output_raises_if_no_matching_call(self):
|
||||
with pytest.raises(ValueError, match="No call message found for"):
|
||||
response_input_to_harmony(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "no_such_id",
|
||||
"output": "x",
|
||||
},
|
||||
prev_responses=[_PREV_CALL],
|
||||
)
|
||||
|
||||
def test_function_call_output_raises_on_empty_prev_responses(self):
|
||||
with pytest.raises(ValueError, match="No call message found for"):
|
||||
response_input_to_harmony(
|
||||
{"type": "function_call_output", "call_id": "call_test", "output": "x"},
|
||||
prev_responses=[],
|
||||
)
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Error cases
|
||||
# -----------------------------------------------------------------------
|
||||
|
||||
def test_unknown_type_raises_value_error(self):
|
||||
with pytest.raises(ValueError, match="Unknown input type"):
|
||||
response_input_to_harmony(
|
||||
{"type": "image_url", "url": "https://example.com/img.png"},
|
||||
prev_responses=[],
|
||||
)
|
||||
+6
-1
@@ -10,7 +10,7 @@ import pybase64 as base64
|
||||
import pytest
|
||||
import websockets
|
||||
|
||||
from tests.entrypoints.openai.conftest import add_attention_backend
|
||||
from tests.entrypoints.speech_to_text.conftest import add_attention_backend
|
||||
from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
@@ -163,6 +163,11 @@ async def test_multi_chunk_streaming(
|
||||
" A little piece of practical poetry. Mary had a little lamb,"
|
||||
" it sleeps with quite a flow, and everywhere that Mary went,"
|
||||
" the lamb was sure to go."
|
||||
) or full_text == (
|
||||
" First words I spoke in the original phonograph."
|
||||
" A little piece of practical poetry. Mary had a little lamb,"
|
||||
" it squeaked with quite a flow, and everywhere that Mary went,"
|
||||
" the lamb was sure to go."
|
||||
)
|
||||
|
||||
|
||||
+2
-2
@@ -7,8 +7,8 @@ from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.openai.speech_to_text.protocol import TranscriptionResponse
|
||||
from vllm.entrypoints.openai.speech_to_text.speech_to_text import OpenAISpeechToText
|
||||
from vllm.entrypoints.speech_to_text.base.serving import OpenAISpeechToText
|
||||
from vllm.entrypoints.speech_to_text.transcription.protocol import TranscriptionResponse
|
||||
|
||||
|
||||
async def _never_finishes():
|
||||
@@ -0,0 +1,61 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def transcription_server_with_force_include_usage():
|
||||
args = [
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-num-seqs",
|
||||
"4",
|
||||
"--enforce-eager",
|
||||
"--enable-force-include-usage",
|
||||
"--gpu-memory-utilization",
|
||||
"0.2",
|
||||
]
|
||||
|
||||
with RemoteOpenAIServer("openai/whisper-large-v3-turbo", args) as remote_server:
|
||||
yield remote_server
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def transcription_client_with_force_include_usage(
|
||||
transcription_server_with_force_include_usage,
|
||||
):
|
||||
async with (
|
||||
transcription_server_with_force_include_usage.get_async_client() as async_client
|
||||
):
|
||||
yield async_client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_transcription_with_enable_force_include_usage(
|
||||
transcription_client_with_force_include_usage, winning_call
|
||||
):
|
||||
res = (
|
||||
await transcription_client_with_force_include_usage.audio.transcriptions.create(
|
||||
model="openai/whisper-large-v3-turbo",
|
||||
file=winning_call,
|
||||
language="en",
|
||||
temperature=0.0,
|
||||
stream=True,
|
||||
timeout=30,
|
||||
)
|
||||
)
|
||||
|
||||
async for chunk in res:
|
||||
if not len(chunk.choices):
|
||||
# final usage sent
|
||||
usage = chunk.usage
|
||||
assert isinstance(usage, dict)
|
||||
assert usage["prompt_tokens"] > 0
|
||||
assert usage["completion_tokens"] > 0
|
||||
assert usage["total_tokens"] > 0
|
||||
else:
|
||||
assert not hasattr(chunk, "usage")
|
||||
+5
-3
@@ -24,12 +24,14 @@ from vllm.entrypoints.openai.engine.protocol import (
|
||||
RequestResponseMetadata,
|
||||
)
|
||||
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
|
||||
from vllm.entrypoints.openai.speech_to_text.protocol import TranscriptionRequest
|
||||
from vllm.entrypoints.openai.speech_to_text.serving import OpenAIServingTranscription
|
||||
from vllm.entrypoints.openai.speech_to_text.speech_to_text import (
|
||||
from vllm.entrypoints.speech_to_text.base.serving import (
|
||||
OpenAISpeechToText,
|
||||
asr_inter_chunk_separator,
|
||||
)
|
||||
from vllm.entrypoints.speech_to_text.transcription.protocol import TranscriptionRequest
|
||||
from vllm.entrypoints.speech_to_text.transcription.serving import (
|
||||
OpenAIServingTranscription,
|
||||
)
|
||||
from vllm.model_executor.models.interfaces import SupportsTranscription
|
||||
from vllm.outputs import CompletionOutput, RequestOutput
|
||||
|
||||
+1
-1
@@ -6,7 +6,7 @@ import json
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.openai.conftest import add_attention_backend
|
||||
from tests.entrypoints.speech_to_text.conftest import add_attention_backend
|
||||
from tests.utils import ROCM_ENV_OVERRIDES, ROCM_EXTRA_ARGS, RemoteOpenAIServer
|
||||
|
||||
MISTRAL_FORMAT_ARGS = [
|
||||
+1
-1
@@ -13,7 +13,7 @@ import pytest
|
||||
import pytest_asyncio
|
||||
import soundfile as sf
|
||||
|
||||
from tests.entrypoints.openai.conftest import add_attention_backend
|
||||
from tests.entrypoints.speech_to_text.conftest import add_attention_backend
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm.logger import init_logger
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
@@ -0,0 +1,3 @@
|
||||
Qwen3-Next-80B-A3B-NVFP4-EP2.yaml
|
||||
Qwen3-Next-FP8-EP2.yaml
|
||||
Nemotron-3-Super-120B-A12B-NVFP4.yaml
|
||||
@@ -3,6 +3,3 @@ Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
|
||||
Qwen1.5-MoE-W4A16-CT.yaml
|
||||
DeepSeek-V2-Lite-Instruct-FP8.yaml
|
||||
Qwen3-30B-A3B-NVFP4.yaml
|
||||
Qwen3-Next-80B-A3B-NVFP4-EP2.yaml
|
||||
Qwen3-Next-FP8-EP2.yaml
|
||||
Nemotron-3-Super-120B-A12B-NVFP4.yaml
|
||||
|
||||
@@ -0,0 +1,377 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_rocm(), reason="Only used by ROCm"
|
||||
)
|
||||
|
||||
NOPE_HEAD_DIM = 448
|
||||
ROPE_HEAD_DIM = 64
|
||||
HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM
|
||||
|
||||
|
||||
def _ref_global_topk_ragged(
|
||||
topk_indices: torch.Tensor,
|
||||
token_to_req_indices: torch.Tensor,
|
||||
block_table: torch.Tensor,
|
||||
block_size: int,
|
||||
is_valid_token: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
topk = topk_indices.reshape(topk_indices.shape[0], -1)
|
||||
valid = (topk >= 0) & is_valid_token[:, None]
|
||||
lens = valid.sum(dim=1, dtype=torch.int32)
|
||||
indptr = torch.zeros(lens.shape[0] + 1, dtype=torch.int32, device=topk.device)
|
||||
torch.cumsum(lens, dim=0, out=indptr[1:])
|
||||
|
||||
safe_topk = torch.clamp(topk, min=0)
|
||||
block_indices = safe_topk // block_size
|
||||
block_offsets = safe_topk % block_size
|
||||
req_indices = token_to_req_indices[:, None].expand_as(topk)
|
||||
slot_ids = block_table[req_indices, block_indices] * block_size + block_offsets
|
||||
|
||||
offsets = torch.arange(topk.shape[1], dtype=torch.int32, device=topk.device)
|
||||
positions = indptr[:-1, None] + offsets[None, :]
|
||||
return slot_ids[valid], positions[valid].to(torch.long), indptr, lens
|
||||
|
||||
|
||||
def _ref_sparse_prefill_ragged(
|
||||
q: torch.Tensor,
|
||||
kv: torch.Tensor,
|
||||
rows: list[list[int]],
|
||||
scale: float,
|
||||
attn_sink: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
q_f32 = q.float()
|
||||
kv_f32 = kv.float()
|
||||
out = torch.empty_like(q_f32)
|
||||
|
||||
for query_idx in range(q.shape[0]):
|
||||
row_indices = rows[query_idx]
|
||||
for head_idx in range(q.shape[1]):
|
||||
if row_indices:
|
||||
selected_kv = kv_f32[row_indices]
|
||||
scores = torch.mv(selected_kv, q_f32[query_idx, head_idx]) * scale
|
||||
if attn_sink is not None:
|
||||
scores_with_sink = torch.cat(
|
||||
[scores, attn_sink[head_idx].float().reshape(1)]
|
||||
)
|
||||
probs = torch.softmax(scores_with_sink, dim=0)[:-1]
|
||||
else:
|
||||
probs = torch.softmax(scores, dim=0)
|
||||
out[query_idx, head_idx] = torch.sum(
|
||||
probs[:, None] * selected_kv, dim=0
|
||||
)
|
||||
else:
|
||||
out[query_idx, head_idx] = 0
|
||||
return out.to(torch.bfloat16)
|
||||
|
||||
|
||||
def _pack_fp8_ds_mla_cache(kv: torch.Tensor, block_size: int) -> torch.Tensor:
|
||||
assert kv.shape[-1] == HEAD_DIM
|
||||
num_tokens = kv.shape[0]
|
||||
num_blocks = (num_tokens + block_size - 1) // block_size
|
||||
cache = torch.zeros(
|
||||
(num_blocks, block_size, 584),
|
||||
dtype=torch.uint8,
|
||||
device=kv.device,
|
||||
)
|
||||
cache_flat = cache.view(torch.uint8).flatten()
|
||||
kv_nope_fp8 = (
|
||||
kv[:, :NOPE_HEAD_DIM].to(current_platform.fp8_dtype()).view(torch.uint8)
|
||||
)
|
||||
kv_rope_u8 = kv[:, NOPE_HEAD_DIM:].contiguous().view(torch.uint8)
|
||||
|
||||
for slot in range(num_tokens):
|
||||
block_idx = slot // block_size
|
||||
pos = slot % block_size
|
||||
block_base = block_idx * cache.stride(0)
|
||||
token_base = block_base + pos * 576
|
||||
scale_base = block_base + block_size * 576 + pos * 8
|
||||
cache_flat[token_base : token_base + NOPE_HEAD_DIM].copy_(kv_nope_fp8[slot])
|
||||
cache_flat[
|
||||
token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
|
||||
].copy_(kv_rope_u8[slot])
|
||||
cache_flat[scale_base : scale_base + 7].fill_(127)
|
||||
return cache
|
||||
|
||||
|
||||
def _read_fp8_ds_mla_cache(
|
||||
cache: torch.Tensor, slot: int, block_size: int
|
||||
) -> torch.Tensor:
|
||||
cache_flat = cache.view(torch.uint8).flatten()
|
||||
block_idx = slot // block_size
|
||||
pos = slot % block_size
|
||||
block_base = block_idx * cache.stride(0)
|
||||
token_base = block_base + pos * 576
|
||||
|
||||
nope_u8 = cache_flat[token_base : token_base + NOPE_HEAD_DIM]
|
||||
nope = nope_u8.view(current_platform.fp8_dtype()).to(torch.float32)
|
||||
rope_u8 = cache_flat[
|
||||
token_base + NOPE_HEAD_DIM : token_base + NOPE_HEAD_DIM + ROPE_HEAD_DIM * 2
|
||||
]
|
||||
rope = rope_u8.view(torch.bfloat16).to(torch.float32)
|
||||
return torch.cat([nope, rope])
|
||||
|
||||
|
||||
def _ref_sparse_decode_ragged(
|
||||
q: torch.Tensor,
|
||||
main_cache: torch.Tensor,
|
||||
main_rows: list[list[int]],
|
||||
scale: float,
|
||||
attn_sink: torch.Tensor | None,
|
||||
block_size: int,
|
||||
extra_cache: torch.Tensor | None = None,
|
||||
extra_rows: list[list[int]] | None = None,
|
||||
) -> torch.Tensor:
|
||||
q_f32 = q.float()
|
||||
out = torch.empty_like(q_f32)
|
||||
|
||||
for query_idx in range(q.shape[0]):
|
||||
row_kv = [
|
||||
_read_fp8_ds_mla_cache(main_cache, int(slot), block_size)
|
||||
for slot in main_rows[query_idx]
|
||||
]
|
||||
if extra_cache is not None and extra_rows is not None:
|
||||
row_kv.extend(
|
||||
_read_fp8_ds_mla_cache(extra_cache, int(slot), block_size)
|
||||
for slot in extra_rows[query_idx]
|
||||
)
|
||||
|
||||
kv = torch.stack(row_kv).to(q.device)
|
||||
for head_idx in range(q.shape[1]):
|
||||
scores = torch.mv(kv, q_f32[query_idx, head_idx]) * scale
|
||||
if attn_sink is not None:
|
||||
scores_with_sink = torch.cat(
|
||||
[scores, attn_sink[head_idx].float().reshape(1)]
|
||||
)
|
||||
probs = torch.softmax(scores_with_sink, dim=0)[:-1]
|
||||
else:
|
||||
probs = torch.softmax(scores, dim=0)
|
||||
out[query_idx, head_idx] = torch.sum(probs[:, None] * kv, dim=0)
|
||||
return out.to(torch.bfloat16)
|
||||
|
||||
|
||||
def _ref_combine_topk_swa_ragged(
|
||||
device: torch.device,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
expected_ragged = torch.tensor(
|
||||
[
|
||||
100,
|
||||
101,
|
||||
7,
|
||||
8,
|
||||
9,
|
||||
110,
|
||||
111,
|
||||
8,
|
||||
9,
|
||||
10,
|
||||
120,
|
||||
121,
|
||||
122,
|
||||
9,
|
||||
10,
|
||||
11,
|
||||
150,
|
||||
27,
|
||||
28,
|
||||
29,
|
||||
160,
|
||||
161,
|
||||
28,
|
||||
29,
|
||||
30,
|
||||
],
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
expected_lens = torch.tensor([5, 5, 6, 4, 5], dtype=torch.int32, device=device)
|
||||
expected_indptr = torch.zeros(6, dtype=torch.int32, device=device)
|
||||
torch.cumsum(expected_lens, dim=0, out=expected_indptr[1:])
|
||||
return expected_ragged, expected_indptr, expected_lens
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_compute_global_topk_ragged_indices_and_indptr() -> None:
|
||||
from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
|
||||
compute_global_topk_ragged_indices_and_indptr,
|
||||
)
|
||||
|
||||
device = torch.device("cuda")
|
||||
block_size = 4
|
||||
topk_indices = torch.tensor(
|
||||
[
|
||||
[0, 3, 4, -1],
|
||||
[5, 8, -1, -1],
|
||||
[2, 7, 9, -1],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
token_to_req_indices = torch.tensor([0, 1, 1], dtype=torch.int32, device=device)
|
||||
block_table = torch.tensor(
|
||||
[
|
||||
[10, 11, 12],
|
||||
[20, 21, 22],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
is_valid_token = torch.tensor([True, False, True], dtype=torch.bool, device=device)
|
||||
|
||||
actual_ragged, actual_indptr, actual_lens = (
|
||||
compute_global_topk_ragged_indices_and_indptr(
|
||||
topk_indices,
|
||||
token_to_req_indices,
|
||||
block_table,
|
||||
block_size,
|
||||
is_valid_token,
|
||||
)
|
||||
)
|
||||
expected_values, expected_positions, expected_indptr, expected_lens = (
|
||||
_ref_global_topk_ragged(
|
||||
topk_indices,
|
||||
token_to_req_indices,
|
||||
block_table,
|
||||
block_size,
|
||||
is_valid_token,
|
||||
)
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual_ragged[expected_positions], expected_values)
|
||||
torch.testing.assert_close(actual_indptr, expected_indptr)
|
||||
torch.testing.assert_close(actual_lens, expected_lens)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_sparse_attn_prefill_ragged_kernel() -> None:
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
||||
_rocm_sparse_attn_prefill_ragged_triton,
|
||||
)
|
||||
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(0)
|
||||
q = torch.randn(3, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
indices = torch.tensor([0, 2, 1, 3, 4], dtype=torch.int32, device=device)
|
||||
indptr = torch.tensor([0, 2, 5, 5], dtype=torch.int32, device=device)
|
||||
attn_sink = torch.tensor([-0.25, 0.0, 0.25], dtype=torch.float32, device=device)
|
||||
scale = HEAD_DIM**-0.5
|
||||
|
||||
actual = _rocm_sparse_attn_prefill_ragged_triton(
|
||||
q=q,
|
||||
kv=kv,
|
||||
indices=indices,
|
||||
indptr=indptr,
|
||||
scale=scale,
|
||||
attn_sink=attn_sink,
|
||||
nope_head_dim=NOPE_HEAD_DIM,
|
||||
rope_head_dim=ROPE_HEAD_DIM,
|
||||
)
|
||||
expected = _ref_sparse_prefill_ragged(
|
||||
q, kv, [[0, 2], [1, 3, 4], []], scale, attn_sink
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_sparse_attn_decode_ragged_kernel() -> None:
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
||||
_rocm_sparse_attn_decode_ragged_triton,
|
||||
)
|
||||
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(1)
|
||||
block_size = 4
|
||||
q = torch.randn(2, 3, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
main_kv = torch.randn(6, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
extra_kv = torch.randn(5, HEAD_DIM, dtype=torch.bfloat16, device=device) * 0.125
|
||||
main_cache = _pack_fp8_ds_mla_cache(main_kv, block_size)
|
||||
extra_cache = _pack_fp8_ds_mla_cache(extra_kv, block_size)
|
||||
main_indices = torch.tensor([0, 2, 4, 1], dtype=torch.int32, device=device)
|
||||
main_indptr = torch.tensor([0, 2, 4], dtype=torch.int32, device=device)
|
||||
extra_indices = torch.tensor([1, 3, 0], dtype=torch.int32, device=device)
|
||||
extra_indptr = torch.tensor([0, 1, 3], dtype=torch.int32, device=device)
|
||||
attn_sink = torch.tensor([-0.1, 0.0, 0.1], dtype=torch.float32, device=device)
|
||||
scale = HEAD_DIM**-0.5
|
||||
|
||||
actual = _rocm_sparse_attn_decode_ragged_triton(
|
||||
q=q,
|
||||
main_cache=main_cache,
|
||||
main_indices=main_indices,
|
||||
main_indptr=main_indptr,
|
||||
scale=scale,
|
||||
attn_sink=attn_sink,
|
||||
nope_head_dim=NOPE_HEAD_DIM,
|
||||
rope_head_dim=ROPE_HEAD_DIM,
|
||||
extra_cache=extra_cache,
|
||||
extra_indices=extra_indices,
|
||||
extra_indptr=extra_indptr,
|
||||
)
|
||||
expected = _ref_sparse_decode_ragged(
|
||||
q=q,
|
||||
main_cache=main_cache,
|
||||
main_rows=[[0, 2], [4, 1]],
|
||||
scale=scale,
|
||||
attn_sink=attn_sink,
|
||||
block_size=block_size,
|
||||
extra_cache=extra_cache,
|
||||
extra_rows=[[1], [3, 0]],
|
||||
)
|
||||
|
||||
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_combine_topk_swa_indices_ragged() -> None:
|
||||
from vllm.v1.attention.backends.mla.rocm_aiter_mla_sparse_dsv4 import (
|
||||
combine_topk_swa_indices_ragged,
|
||||
)
|
||||
|
||||
device = torch.device("cuda")
|
||||
topk_indices = torch.tensor(
|
||||
[
|
||||
[100, 101, 102, 103],
|
||||
[110, 111, 112, 113],
|
||||
[120, 121, 122, 123],
|
||||
[130, 131, 132, 133],
|
||||
[140, 141, 142, 143],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
|
||||
seq_lens = torch.tensor([6, 4], dtype=torch.int32, device=device)
|
||||
gather_lens = torch.tensor([4, 3], dtype=torch.int32, device=device)
|
||||
window_size = 3
|
||||
compress_ratio = 2
|
||||
topk = 4
|
||||
M = 20
|
||||
N = 8
|
||||
|
||||
actual_ragged, actual_indptr, actual_lens = combine_topk_swa_indices_ragged(
|
||||
topk_indices,
|
||||
query_start_loc,
|
||||
seq_lens,
|
||||
gather_lens,
|
||||
window_size,
|
||||
compress_ratio,
|
||||
topk,
|
||||
M,
|
||||
N,
|
||||
)
|
||||
expected_ragged, expected_indptr, expected_lens = _ref_combine_topk_swa_ragged(
|
||||
device
|
||||
)
|
||||
|
||||
torch.testing.assert_close(
|
||||
actual_ragged[: expected_ragged.numel()], expected_ragged
|
||||
)
|
||||
torch.testing.assert_close(actual_indptr, expected_indptr)
|
||||
torch.testing.assert_close(actual_lens, expected_lens)
|
||||
@@ -66,24 +66,6 @@ def reference(x: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Ten
|
||||
return reference_quant(ref_act_out, use_ue8m0)
|
||||
|
||||
|
||||
def reference_with_clamp(
|
||||
x: torch.Tensor, use_ue8m0: bool, clamp_limit: float
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Pre-clamp inputs (gate from above, up symmetric) at the input dtype to
|
||||
match the C++ compute() template, then run the standard silu_and_mul +
|
||||
quant reference."""
|
||||
N_2 = x.size(1) // 2
|
||||
dtype = x.dtype
|
||||
gate = x[..., :N_2].to(torch.float32).clamp(max=clamp_limit).to(dtype)
|
||||
up = (
|
||||
x[..., N_2:]
|
||||
.to(torch.float32)
|
||||
.clamp(min=-clamp_limit, max=clamp_limit)
|
||||
.to(dtype)
|
||||
)
|
||||
return reference(torch.cat([gate, up], dim=-1), use_ue8m0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("T", [128, 256, 512])
|
||||
@pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2])
|
||||
@pytest.mark.skipif(
|
||||
@@ -107,32 +89,3 @@ def test_silu_mul_fp8_quant_deep_gemm(T: int, N: int):
|
||||
|
||||
torch.testing.assert_close(output.to(torch.float32), ref_output.to(torch.float32))
|
||||
torch.testing.assert_close(output_scales, ref_output_scales)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("T", [128, 256, 512])
|
||||
@pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2])
|
||||
@pytest.mark.parametrize("clamp_limit", [7.0, 10.0])
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_rocm(),
|
||||
reason="ROCm does not support DeepGemm.",
|
||||
)
|
||||
def test_silu_mul_fp8_quant_deep_gemm_clamp(T: int, N: int, clamp_limit: float):
|
||||
set_random_seed(42)
|
||||
|
||||
# Use a wide distribution so values routinely exceed both clamp limits and
|
||||
# the clamp branch is actually exercised (uniform [0, 1) inputs would never
|
||||
# trigger it).
|
||||
input = torch.randn((T, N), dtype=torch.bfloat16, device="cuda") * 8.0
|
||||
|
||||
use_ue8m0 = is_deep_gemm_e8m0_used()
|
||||
|
||||
# Test
|
||||
output, output_scales = silu_mul_per_token_group_quant_fp8_colmajor(
|
||||
input, use_ue8m0=use_ue8m0, clamp_limit=clamp_limit
|
||||
)
|
||||
|
||||
# Reference
|
||||
ref_output, ref_output_scales = reference_with_clamp(input, use_ue8m0, clamp_limit)
|
||||
|
||||
torch.testing.assert_close(output.to(torch.float32), ref_output.to(torch.float32))
|
||||
torch.testing.assert_close(output_scales, ref_output_scales)
|
||||
|
||||
@@ -3,11 +3,12 @@
|
||||
"""
|
||||
Round-trip tests for compressor → FP8 quant + KV cache insert → gather + dequant.
|
||||
|
||||
Four test functions cover five paths:
|
||||
These tests cover:
|
||||
A) DeepseekV4 Attention: head_dim=512 (448 FP8 nope + 64 bf16 rope), quant_block=64
|
||||
B) Indexer: head_dim=128 (all FP8), quant_block=128
|
||||
C) DeepseekV4 Attention magnitude range: correctness across small/large values
|
||||
D) Indexer fused Triton kernel: compress+norm+rope+quant+insert
|
||||
B) Fused dequant+gather K cache
|
||||
C) Indexer: head_dim=128 (all FP8), quant_block=128
|
||||
D) DeepseekV4 Attention magnitude range: correctness across small/large values
|
||||
E) Indexer fused Triton kernel: compress+norm+rope+quant+insert
|
||||
"""
|
||||
|
||||
import math
|
||||
@@ -134,7 +135,140 @@ def test_deepseek_v4_attention_quant_cache_roundtrip(num_tokens: int, block_size
|
||||
)
|
||||
|
||||
|
||||
# ── Test B: Indexer path ────────────────────────────────────────────────────
|
||||
# ── Test B: Fused dequant+gather K cache ────────────────────────────────────
|
||||
|
||||
|
||||
def _dequantize_and_gather_k_cache_reference(
|
||||
out: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
gather_lens: torch.Tensor | None,
|
||||
block_table: torch.Tensor,
|
||||
block_size: int,
|
||||
offset: int,
|
||||
) -> None:
|
||||
fp8_dim = 448
|
||||
bf16_dim = 64
|
||||
scale_dim = 8
|
||||
quant_block = 64
|
||||
token_data_size = fp8_dim + bf16_dim * 2
|
||||
|
||||
for req_id in range(seq_lens.shape[0]):
|
||||
seq_len = seq_lens[req_id].item()
|
||||
gather_len = gather_lens[req_id].item() if gather_lens is not None else seq_len
|
||||
start_pos = seq_len - gather_len
|
||||
|
||||
for i in range(gather_len):
|
||||
pos = start_pos + i
|
||||
pos_in_block = pos % block_size
|
||||
block_idx = block_table[req_id, pos // block_size].item()
|
||||
cache_block = k_cache[block_idx].view(-1)
|
||||
|
||||
token_data_start = pos_in_block * token_data_size
|
||||
fp8_bytes = cache_block[token_data_start : token_data_start + fp8_dim]
|
||||
fp8_vals = fp8_bytes.view(torch.float8_e4m3fn).float()
|
||||
|
||||
scale_start = block_size * token_data_size + pos_in_block * scale_dim
|
||||
encoded_scales = cache_block[scale_start : scale_start + scale_dim]
|
||||
scales = torch.exp2(encoded_scales[:7].float() - 127.0)
|
||||
dequant = fp8_vals * scales.repeat_interleave(quant_block)
|
||||
|
||||
bf16_start = token_data_start + fp8_dim
|
||||
bf16_bytes = cache_block[bf16_start : bf16_start + bf16_dim * 2]
|
||||
bf16_tail = bf16_bytes.view(torch.bfloat16)
|
||||
|
||||
out[req_id, offset + i, :fp8_dim] = dequant
|
||||
out[req_id, offset + i, fp8_dim:] = bf16_tail
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("seq_lens_host", "gather_lens_host", "offset"),
|
||||
[
|
||||
([9, 23, 7], None, 0),
|
||||
([19, 8, 257], [6, 8, 129], 5),
|
||||
],
|
||||
)
|
||||
def test_dequantize_and_gather_k_cache(
|
||||
seq_lens_host: list[int],
|
||||
gather_lens_host: list[int] | None,
|
||||
offset: int,
|
||||
):
|
||||
block_size = 64
|
||||
head_dim = 512
|
||||
nope_dim = 448
|
||||
scale_dim = 8
|
||||
head_bytes = nope_dim + (head_dim - nope_dim) * 2 + scale_dim
|
||||
device = "cuda"
|
||||
num_reqs = len(seq_lens_host)
|
||||
num_tokens = sum(seq_lens_host)
|
||||
max_gather_len = max(gather_lens_host or seq_lens_host)
|
||||
max_blocks_per_seq = math.ceil(max(seq_lens_host) / block_size)
|
||||
num_blocks = sum(math.ceil(seq_len / block_size) for seq_len in seq_lens_host)
|
||||
|
||||
compressed_kv = torch.randn(
|
||||
num_tokens, head_dim, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
|
||||
# Randomize physical pages so the test covers block-table translation.
|
||||
# Keep padded block-table entries invalid to catch accidental reads.
|
||||
physical_blocks = torch.randperm(num_blocks, device=device)
|
||||
block_table = torch.full(
|
||||
(num_reqs, max_blocks_per_seq), int(-1e6), dtype=torch.int32, device=device
|
||||
)
|
||||
start = 0
|
||||
for req_id, seq_len in enumerate(seq_lens_host):
|
||||
num_req_blocks = math.ceil(seq_len / block_size)
|
||||
req_blocks = physical_blocks[start : start + num_req_blocks]
|
||||
block_table[req_id, :num_req_blocks] = req_blocks
|
||||
start += num_req_blocks
|
||||
|
||||
# Build slot_mapping for quantize_and_insert_k_cache.
|
||||
slot_mapping = torch.empty(num_tokens, dtype=torch.int64, device=device)
|
||||
start = 0
|
||||
for req_id, seq_len in enumerate(seq_lens_host):
|
||||
logical_pos = torch.arange(seq_len, dtype=torch.int64, device=device)
|
||||
block_idx = block_table[req_id, logical_pos // block_size].to(torch.int64)
|
||||
token_slots = block_idx * block_size + logical_pos % block_size
|
||||
slot_mapping[start : start + seq_len] = token_slots
|
||||
start += seq_len
|
||||
|
||||
# Insert compressed K into the paged cache layout used by the gather op.
|
||||
k_cache = torch.empty(
|
||||
num_blocks, block_size, head_bytes, dtype=torch.uint8, device=device
|
||||
)
|
||||
k_cache_2d = k_cache.view(num_blocks, -1)
|
||||
quantize_and_insert_k_cache(compressed_kv, k_cache_2d, slot_mapping, block_size)
|
||||
|
||||
out_shape = (num_reqs, offset + max_gather_len + 3, head_dim)
|
||||
ref_out = torch.empty(out_shape, dtype=torch.bfloat16, device=device)
|
||||
actual_out = torch.empty_like(ref_out)
|
||||
seq_lens = torch.tensor(seq_lens_host, dtype=torch.int32, device=device)
|
||||
gather_lens = (
|
||||
torch.tensor(gather_lens_host, dtype=torch.int32, device=device)
|
||||
if gather_lens_host is not None
|
||||
else None
|
||||
)
|
||||
|
||||
# Compare production gather against a PyTorch reference for valid output rows.
|
||||
_dequantize_and_gather_k_cache_reference(
|
||||
ref_out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
|
||||
)
|
||||
dequantize_and_gather_k_cache(
|
||||
actual_out, k_cache, seq_lens, gather_lens, block_table, block_size, offset
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
# only check non-padded content
|
||||
for req_id, seq_len in enumerate(seq_lens_host):
|
||||
gather_len = (
|
||||
gather_lens_host[req_id] if gather_lens_host is not None else seq_len
|
||||
)
|
||||
actual = actual_out[req_id, offset : offset + gather_len]
|
||||
expected = ref_out[req_id, offset : offset + gather_len]
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
# ── Test C: Indexer path ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", [1, 4, 8, 17])
|
||||
@@ -254,7 +388,7 @@ def test_indexer_gather_accepts_upper_bound_output():
|
||||
assert torch.all(dst_scale[valid_tokens:] == sentinel)
|
||||
|
||||
|
||||
# ── Test C: DeepseekV4 attention with values at different magnitudes ───────────
|
||||
# ── Test D: DeepseekV4 attention with values at different magnitudes ───────────
|
||||
|
||||
|
||||
def test_deepseek_v4_quant_magnitude_range():
|
||||
@@ -316,7 +450,7 @@ def test_deepseek_v4_quant_magnitude_range():
|
||||
)
|
||||
|
||||
|
||||
# ── Test D: Indexer fused K-cache insert (Triton kernels) ────────────────────
|
||||
# ── Test E: Indexer fused K-cache insert (Triton kernels) ────────────────────
|
||||
#
|
||||
# Both kernels share the same Triton signature; use_fp4 selects between them.
|
||||
# Full pipeline: state-cache gather → softmax-weighted compress → RMSNorm →
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
import vllm
|
||||
import vllm.config
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ..utils import create_new_process_for_each_test, multi_gpu_test
|
||||
|
||||
@@ -50,6 +53,9 @@ def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
|
||||
return generated_texts
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
@create_new_process_for_each_test()
|
||||
def test_chatglm3_lora(chatglm3_lora_files):
|
||||
llm = vllm.LLM(
|
||||
@@ -70,6 +76,9 @@ def test_chatglm3_lora(chatglm3_lora_files):
|
||||
assert output2[i] == EXPECTED_LORA_OUTPUT[i]
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
def test_chatglm3_lora_tp4(chatglm3_lora_files):
|
||||
llm = vllm.LLM(
|
||||
|
||||
@@ -11,6 +11,7 @@ import pytest
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ..conftest import AudioTestAssets, VllmRunner
|
||||
from ..utils import create_new_process_for_each_test
|
||||
@@ -76,6 +77,9 @@ def test_active_default_mm_lora(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
@create_new_process_for_each_test()
|
||||
def test_inactive_default_mm_lora(
|
||||
vllm_runner: type[VllmRunner],
|
||||
@@ -92,6 +96,9 @@ def test_inactive_default_mm_lora(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
@create_new_process_for_each_test()
|
||||
def test_default_mm_lora_succeeds_with_redundant_lora_request(
|
||||
vllm_runner: type[VllmRunner],
|
||||
@@ -107,6 +114,9 @@ def test_default_mm_lora_succeeds_with_redundant_lora_request(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
@create_new_process_for_each_test()
|
||||
def test_default_mm_lora_fails_with_overridden_lora_request(
|
||||
vllm_runner: type[VllmRunner],
|
||||
|
||||
@@ -10,6 +10,7 @@ import vllm.config
|
||||
from vllm import LLM
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ..utils import VLLM_PATH, create_new_process_for_each_test, multi_gpu_test
|
||||
|
||||
@@ -139,6 +140,9 @@ def test_llama_lora(llama32_lora_files, cudagraph_specialize_lora: bool):
|
||||
generate_and_test(llm, llama32_lora_files)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
def test_llama_lora_tp4(llama32_lora_files):
|
||||
llm = vllm.LLM(
|
||||
|
||||
@@ -68,6 +68,9 @@ def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
|
||||
return generated_texts
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
def test_minicpmv_lora(minicpmv_lora_files):
|
||||
llm = vllm.LLM(
|
||||
MODEL_PATH,
|
||||
|
||||
@@ -11,6 +11,7 @@ from safetensors.torch import load_file, save_file
|
||||
|
||||
import vllm
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ..utils import multi_gpu_test
|
||||
|
||||
@@ -110,6 +111,9 @@ def generate_and_test(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
def test_olmoe_lora(olmoe_lora_files, maybe_enable_lora_dual_stream):
|
||||
# We enable enforce_eager=True here to reduce VRAM usage for lora-test CI,
|
||||
# Otherwise, the lora-test will fail due to CUDA OOM.
|
||||
@@ -178,6 +182,9 @@ def test_olmoe_lora_mixed_random(
|
||||
assert outputs[0].outputs[0].text.strip().startswith(EXPECTED_LORA_OUTPUT[0])
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
@pytest.mark.parametrize("fully_sharded_loras", [False, True])
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
def test_olmoe_lora_tp2(olmoe_lora_files, fully_sharded_loras):
|
||||
|
||||
@@ -1,12 +1,14 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
import vllm
|
||||
import vllm.config
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ..utils import create_new_process_for_each_test, multi_gpu_test
|
||||
|
||||
@@ -311,6 +313,9 @@ def _assert_qwen35_text_vl_and_mixed_lora(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
@create_new_process_for_each_test()
|
||||
def test_qwen35_text_lora(
|
||||
qwen35_text_lora_files, qwen35_vl_lora_files, maybe_enable_lora_dual_stream
|
||||
|
||||
@@ -2,12 +2,14 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
from packaging.version import Version
|
||||
from transformers import __version__ as TRANSFORMERS_VERSION
|
||||
|
||||
import vllm
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.sampling_params import BeamSearchParams
|
||||
|
||||
|
||||
@@ -206,6 +208,9 @@ def test_qwen2vl_lora_beam_search(qwen2vl_lora_files):
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
def test_qwen25vl_lora(qwen25vl_lora_files):
|
||||
"""Test Qwen 2.5 VL model with LoRA"""
|
||||
config = TestConfig(model_path=QWEN25VL_MODEL_PATH, lora_path=qwen25vl_lora_files)
|
||||
@@ -216,6 +221,9 @@ def test_qwen25vl_lora(qwen25vl_lora_files):
|
||||
tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_cuda_alike(), reason="Skipping to avoid redundant model tests"
|
||||
)
|
||||
def test_qwen25vl_vision_lora(qwen25vl_vision_lora_files):
|
||||
config = TestConfig(
|
||||
model_path=QWEN25VL_MODEL_PATH,
|
||||
|
||||
@@ -124,30 +124,3 @@ def test_whisper_multi_lora(whisper_lora_files):
|
||||
f"Expected same outputs for same adapter with different IDs. "
|
||||
f"Got: {outputs_lora1} vs {outputs_lora2}"
|
||||
)
|
||||
|
||||
|
||||
@create_new_process_for_each_test()
|
||||
def test_whisper_with_and_without_lora(whisper_lora_files):
|
||||
"""Test that Whisper produces different outputs with and without LoRA.
|
||||
|
||||
This test verifies that the LoRA adapter actually affects the model output.
|
||||
"""
|
||||
llm = create_whisper_llm(enable_lora=True)
|
||||
|
||||
# Run with LoRA
|
||||
outputs_with_lora = run_whisper_inference(
|
||||
llm, lora_path=whisper_lora_files, lora_id=1
|
||||
)
|
||||
|
||||
# Run without LoRA (base model only)
|
||||
outputs_without_lora = run_whisper_inference(llm, lora_path=None)
|
||||
|
||||
# Both should produce valid outputs
|
||||
assert len(outputs_with_lora[0]) > 0
|
||||
assert len(outputs_without_lora[0]) > 0
|
||||
|
||||
print(f"Output with LoRA: {outputs_with_lora[0]}")
|
||||
print(f"Output without LoRA: {outputs_without_lora[0]}")
|
||||
|
||||
# Note: Outputs may or may not differ depending on the adapter
|
||||
# The main verification is that both configurations work
|
||||
|
||||
@@ -4,9 +4,12 @@
|
||||
from collections.abc import Mapping
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image as PILImage
|
||||
|
||||
from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs
|
||||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||||
from vllm.multimodal.inputs import MultiModalFieldConfig
|
||||
|
||||
from ....conftest import ImageTestAssets
|
||||
from ...utils import build_model_context
|
||||
@@ -15,6 +18,36 @@ from ...utils import build_model_context
|
||||
GEMMA4_MODEL_ID = "google/gemma-4-E2B-it"
|
||||
|
||||
|
||||
def test_gemma4_image_schema_accepts_variable_patch_counts():
|
||||
Gemma4ImagePixelInputs(
|
||||
pixel_values=[
|
||||
torch.randn(10080, 768),
|
||||
torch.randn(2520, 768),
|
||||
],
|
||||
pixel_position_ids=[
|
||||
torch.zeros(10080, 2, dtype=torch.long),
|
||||
torch.zeros(2520, 2, dtype=torch.long),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_gemma4_image_batching_keeps_variable_patch_counts_unstacked():
|
||||
field = MultiModalFieldConfig.batched("image").field
|
||||
elems = field.build_elems(
|
||||
"image",
|
||||
"pixel_values",
|
||||
[torch.randn(10080, 768), torch.randn(2520, 768)],
|
||||
)
|
||||
|
||||
reduced = field.reduce_data(list(elems))
|
||||
|
||||
assert isinstance(reduced, list)
|
||||
assert [tensor.shape for tensor in reduced] == [
|
||||
torch.Size([10080, 768]),
|
||||
torch.Size([2520, 768]),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"image_width,image_height,max_soft_tokens",
|
||||
[
|
||||
|
||||
@@ -53,6 +53,19 @@ class _VisionModel:
|
||||
self.loaded_weights = list(weights)
|
||||
|
||||
|
||||
class _FakeTensor:
|
||||
"""Sentinel stand-in for torch.Tensor in load_weights tests. Supports the
|
||||
.detach().clone() chain used by load_weights for buffered mm weights;
|
||||
both methods return self so identity (and the existing equality
|
||||
assertions) are preserved through cloning."""
|
||||
|
||||
def detach(self):
|
||||
return self
|
||||
|
||||
def clone(self):
|
||||
return self
|
||||
|
||||
|
||||
def test_nano_nemotron_vl_skips_multimodal_weights_in_text_only_mode():
|
||||
model = object.__new__(NemotronH_Nano_VL_V2)
|
||||
language_model = _LanguageModel()
|
||||
@@ -86,7 +99,7 @@ def test_nano_nemotron_vl_loads_vision_weights_without_sound_encoder():
|
||||
object.__setattr__(model, "sound_encoder", None)
|
||||
|
||||
language_weight = object()
|
||||
vision_weight = object()
|
||||
vision_weight = _FakeTensor()
|
||||
model.load_weights(
|
||||
[
|
||||
("language_model.layers.0.weight", language_weight),
|
||||
|
||||
@@ -28,6 +28,14 @@ def test_getattr_without_cache(monkeypatch: pytest.MonkeyPatch):
|
||||
assert not hasattr(envs.__getattr__, "cache_info")
|
||||
|
||||
|
||||
def test_nixl_side_channel_host_is_not_compile_factor(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setenv("VLLM_NIXL_SIDE_CHANNEL_HOST", "10.0.0.15")
|
||||
|
||||
assert "VLLM_NIXL_SIDE_CHANNEL_HOST" not in envs.compile_factors()
|
||||
|
||||
|
||||
def test_getattr_with_cache(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("VLLM_HOST_IP", "1.1.1.1")
|
||||
monkeypatch.setenv("VLLM_PORT", "1234")
|
||||
|
||||
@@ -6,6 +6,7 @@ import os
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.ray.ray_env import get_env_vars_to_copy
|
||||
from vllm.v1.executor.ray_utils import WORKER_SPECIFIC_ENV_VARS
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Default prefix matching
|
||||
@@ -106,6 +107,19 @@ class TestExclusion:
|
||||
result = get_env_vars_to_copy(exclude_vars={"CUDA_VISIBLE_DEVICES"})
|
||||
assert "CUDA_VISIBLE_DEVICES" not in result
|
||||
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"VLLM_HOST_IP": "10.0.0.1",
|
||||
"VLLM_NIXL_SIDE_CHANNEL_HOST": "10.0.0.1",
|
||||
},
|
||||
clear=False,
|
||||
)
|
||||
def test_worker_specific_host_vars_are_excluded(self):
|
||||
result = get_env_vars_to_copy(exclude_vars=WORKER_SPECIFIC_ENV_VARS)
|
||||
assert "VLLM_HOST_IP" not in result
|
||||
assert "VLLM_NIXL_SIDE_CHANNEL_HOST" not in result
|
||||
|
||||
@patch.dict(os.environ, {"LMCACHE_LOCAL_CPU": "True"}, clear=False)
|
||||
@patch(
|
||||
"vllm.ray.ray_env.RAY_NON_CARRY_OVER_ENV_VARS",
|
||||
|
||||
@@ -20,20 +20,17 @@ from tests.v1.attention.utils import (
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.config.vllm import set_current_vllm_config
|
||||
from vllm.model_executor.layers.attention.mla_attention import (
|
||||
MLAAttention,
|
||||
QueryLenSupport,
|
||||
_DecodeConcatQuantFP8,
|
||||
)
|
||||
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
|
||||
from vllm.v1.attention.backend import CommonAttentionMetadata
|
||||
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
|
||||
from vllm.v1.attention.backends.mla.prefill import (
|
||||
MLAPrefillBackendEnum,
|
||||
get_mla_prefill_backend,
|
||||
)
|
||||
from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
|
||||
from vllm.v1.kv_cache_interface import MLAAttentionSpec
|
||||
@@ -44,7 +41,6 @@ BACKENDS_TO_TEST = [
|
||||
AttentionBackendEnum.FLASH_ATTN_MLA,
|
||||
AttentionBackendEnum.FLASHINFER_MLA,
|
||||
AttentionBackendEnum.TRITON_MLA,
|
||||
AttentionBackendEnum.TOKENSPEED_MLA,
|
||||
]
|
||||
|
||||
DEVICE_TYPE = current_platform.device_type
|
||||
@@ -53,7 +49,6 @@ DEVICE_TYPE = current_platform.device_type
|
||||
if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
|
||||
BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)
|
||||
BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHINFER_MLA)
|
||||
BACKENDS_TO_TEST.remove(AttentionBackendEnum.TOKENSPEED_MLA)
|
||||
|
||||
# Remove FLASH_ATTN_MLA from the list if not supported
|
||||
if not flash_attn_supports_mla():
|
||||
@@ -63,22 +58,6 @@ if not flash_attn_supports_mla():
|
||||
if not is_flashmla_dense_supported()[0]:
|
||||
BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHMLA)
|
||||
|
||||
# Remove TOKENSPEED_MLA if the optional package is not installed
|
||||
if AttentionBackendEnum.TOKENSPEED_MLA in BACKENDS_TO_TEST:
|
||||
try:
|
||||
import tokenspeed_mla # noqa: F401
|
||||
except ImportError:
|
||||
BACKENDS_TO_TEST.remove(AttentionBackendEnum.TOKENSPEED_MLA)
|
||||
|
||||
|
||||
# Filtered per-test via validate_configuration (capability/deps/dims).
|
||||
PREFILL_BACKENDS_TO_TEST = [
|
||||
MLAPrefillBackendEnum.FLASH_ATTN,
|
||||
MLAPrefillBackendEnum.FLASHINFER,
|
||||
MLAPrefillBackendEnum.TRTLLM_RAGGED,
|
||||
MLAPrefillBackendEnum.TOKENSPEED_MLA,
|
||||
]
|
||||
|
||||
|
||||
SPEC_DECODE_BACKENDS = []
|
||||
for backend in BACKENDS_TO_TEST:
|
||||
@@ -410,18 +389,14 @@ class MockSparseMLAAttentionLayer:
|
||||
return output
|
||||
|
||||
|
||||
class MockMLAAttentionLayer(MLAAttention):
|
||||
class MockMLAAttentionLayer(AttentionLayerBase):
|
||||
"""A mock MLA attention layer for testing.
|
||||
|
||||
This replicates the forward_impl logic from MLAAttention to allow
|
||||
testing MLA backends without the full layer infrastructure.
|
||||
|
||||
Subclasses MLAAttention so that backends that filter
|
||||
`static_forward_context` by `isinstance(layer, MLAAttention)` (e.g.
|
||||
FlashInfer prefill, which reads sm_scale through that filter) see the
|
||||
mock as a real MLA layer. MLAAttention.__init__ is intentionally
|
||||
skipped — it would create its own impl/prefill_backend and self-register
|
||||
in static_forward_context, which fights what the test sets up below.
|
||||
The W_UK_T and W_UV weight matrices are created on the layer (like in
|
||||
MLAAttention.process_weights_after_loading), not on the impl.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -437,7 +412,6 @@ class MockMLAAttentionLayer(MLAAttention):
|
||||
q_scale: float,
|
||||
k_scale: float,
|
||||
):
|
||||
torch.nn.Module.__init__(self)
|
||||
self.impl = impl
|
||||
self.num_heads = num_heads
|
||||
self.qk_nope_head_dim = qk_nope_head_dim
|
||||
@@ -588,15 +562,11 @@ def run_attention_backend(
|
||||
q_scale: float,
|
||||
k_scale: float,
|
||||
kv_cache_dtype: str = "auto",
|
||||
prefill_backend: MLAPrefillBackendEnum | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run attention computation using the specified backend's AttentionImpl."""
|
||||
|
||||
builder_cls, impl_cls = try_get_attention_backend(backend)
|
||||
|
||||
# Force the prefill backend selection (None means auto-select).
|
||||
vllm_config.attention_config.mla_prefill_backend = prefill_backend
|
||||
|
||||
# Set the current vllm config so that get_current_vllm_config() works
|
||||
# in the backend implementations
|
||||
with set_current_vllm_config(vllm_config):
|
||||
@@ -608,11 +578,7 @@ def run_attention_backend(
|
||||
vllm_config.parallel_config
|
||||
)
|
||||
head_size = vllm_config.model_config.get_head_size()
|
||||
# Production MLA passes 1/sqrt(qk_head_dim) (the prefill scale) to the
|
||||
# impl and forwards the same value to the prefill backend. FLASHINFER
|
||||
# prefill reads sm_scale back from impl.scale via global_hyperparameters
|
||||
# at plan() time, so impl.scale must agree with prefill_backend.scale.
|
||||
scale = (qk_nope_head_dim + qk_rope_head_dim) ** -0.5
|
||||
scale = 1.0 / (head_size**0.5)
|
||||
impl = impl_cls(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
@@ -717,7 +683,6 @@ def run_attention_backend(
|
||||
@pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16])
|
||||
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
|
||||
@pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])
|
||||
@pytest.mark.parametrize("prefill_backend", PREFILL_BACKENDS_TO_TEST)
|
||||
def test_backend_correctness(
|
||||
default_vllm_config,
|
||||
dist_init,
|
||||
@@ -728,7 +693,6 @@ def test_backend_correctness(
|
||||
kv_cache_dtype: str,
|
||||
q_scale: float,
|
||||
k_scale: float,
|
||||
prefill_backend: MLAPrefillBackendEnum,
|
||||
):
|
||||
"""
|
||||
Test that all backends produce similar outputs to a reference implementation
|
||||
@@ -765,24 +729,6 @@ def test_backend_correctness(
|
||||
if not backends_to_test:
|
||||
pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}")
|
||||
|
||||
# Skip prefill backends that can't satisfy capability/deps/R1 constraints.
|
||||
from vllm.v1.attention.backends.mla.prefill.selector import (
|
||||
MLAPrefillSelectorConfig,
|
||||
)
|
||||
|
||||
try:
|
||||
prefill_invalid_reasons = prefill_backend.get_class().validate_configuration(
|
||||
current_platform.get_device_capability(),
|
||||
MLAPrefillSelectorConfig(dtype=torch.bfloat16, is_r1_compatible=True),
|
||||
)
|
||||
except ImportError:
|
||||
prefill_invalid_reasons = ["ImportError"]
|
||||
if prefill_invalid_reasons:
|
||||
pytest.skip(
|
||||
f"Prefill backend {prefill_backend.name} unavailable: "
|
||||
f"{prefill_invalid_reasons}"
|
||||
)
|
||||
|
||||
batch_spec = BATCH_SPECS[batch_spec_name]
|
||||
is_spec_decode_test = batch_spec_name.startswith("spec_decode")
|
||||
unique_block_sizes = sorted(set(BACKEND_BLOCK_SIZES[b] for b in backends_to_test))
|
||||
@@ -853,13 +799,9 @@ def test_backend_correctness(
|
||||
assert kv_lora_rank + qk_rope_head_dim == head_size, (
|
||||
f"MLA dimensions don't match: {total_head_size} != {head_size}"
|
||||
)
|
||||
decode_scale = 1.0 / (total_head_size**0.5)
|
||||
qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
|
||||
prefill_scale = qk_head_dim**-0.5
|
||||
# MLA reuses prefill_scale for the decode path: production sets
|
||||
# impl.scale = 1/sqrt(qk_head_dim) and the decode kernels apply it even
|
||||
# though the latent attention runs at head_size dimensions. Keeping the
|
||||
# reference here in sync with run_attention_backend's impl.scale.
|
||||
decode_scale = prefill_scale
|
||||
|
||||
# 2. Generate data and compute SDPA reference output for MLA
|
||||
all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
|
||||
@@ -1150,7 +1092,6 @@ def test_backend_correctness(
|
||||
qk_rope_head_dim,
|
||||
v_head_dim,
|
||||
mock_kv_b_proj,
|
||||
prefill_backend=prefill_backend,
|
||||
q_scale=q_scale,
|
||||
k_scale=k_scale,
|
||||
kv_cache_dtype=kv_cache_dtype,
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm.envs as envs
|
||||
from tests.utils import TestFP8Layer, requires_fp8
|
||||
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
|
||||
CutlassFP8ScaledMMLinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
kFp8DynamicTokenSym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
pytest.importorskip("torch.cuda")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup_cuda():
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip("CUTLASS FP8 kernels require CUDA.")
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
|
||||
@requires_fp8
|
||||
@pytest.mark.parametrize("weight_shape", [(1024, 2048), (4608, 4096)])
|
||||
@pytest.mark.parametrize("batch_size", [1, 16, 17, 32, 64, 65, 256, 257])
|
||||
@torch.inference_mode()
|
||||
def test_cutlass_fp8_batch_invariant_fixed_config(
|
||||
weight_shape: tuple[int, int],
|
||||
batch_size: int,
|
||||
default_vllm_config,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
|
||||
monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True)
|
||||
|
||||
torch.manual_seed(0)
|
||||
layer = TestFP8Layer(
|
||||
weight_shape=weight_shape,
|
||||
activation_quant_key=kFp8DynamicTokenSym,
|
||||
weight_quant_key=kFp8StaticTensorSym,
|
||||
input_dtype=torch.bfloat16,
|
||||
out_dtype=torch.bfloat16,
|
||||
device=torch.device("cuda"),
|
||||
force_kernel=CutlassFP8ScaledMMLinearKernel,
|
||||
)
|
||||
assert isinstance(layer.kernel, CutlassFP8ScaledMMLinearKernel)
|
||||
|
||||
in_features = weight_shape[1]
|
||||
needle = torch.randn((1, in_features), device="cuda", dtype=torch.bfloat16)
|
||||
baseline = layer(needle)[0]
|
||||
|
||||
filler = torch.randn(
|
||||
(max(batch_size - 1, 0), in_features), device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
|
||||
front_batch = torch.cat([needle, filler], dim=0)
|
||||
back_batch = torch.cat([filler, needle], dim=0)
|
||||
|
||||
front_output = layer(front_batch)[0]
|
||||
back_output = layer(back_batch)[-1]
|
||||
|
||||
torch.testing.assert_close(front_output, baseline, rtol=0, atol=0)
|
||||
torch.testing.assert_close(back_output, baseline, rtol=0, atol=0)
|
||||
@@ -488,6 +488,10 @@ def _run_eagle_correctness(
|
||||
|
||||
|
||||
@single_gpu_only
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_device_capability_family(100),
|
||||
reason="DeepSeek head_dim=192 not supported on SM100/SM110 (Blackwell)",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
[
|
||||
"model_setup",
|
||||
@@ -718,7 +722,15 @@ def test_eagle_correctness_heavy(
|
||||
["model_setup", "mm_enabled", "expected_accuracy_threshold"],
|
||||
[
|
||||
(("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False, 0.5), # ref: 65%-70%
|
||||
(("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1), False, 0.0), # dummy model
|
||||
pytest.param(
|
||||
("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1),
|
||||
False,
|
||||
0.0,
|
||||
marks=pytest.mark.skipif(
|
||||
current_platform.is_device_capability_family(100),
|
||||
reason="DeepSeek MTP: TRTLLM MoE top_k check fails on Blackwell",
|
||||
),
|
||||
), # dummy model
|
||||
(
|
||||
("mtp", "Qwen/Qwen3.5-0.8B-Base", 1),
|
||||
False,
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import ray
|
||||
|
||||
from vllm.v1.engine.core import EngineCoreActorMixin
|
||||
from vllm.v1.engine.utils import CoreEngineActorManager, EngineZmqAddresses
|
||||
|
||||
|
||||
class _StubEngineCoreActor(EngineCoreActorMixin):
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: Any,
|
||||
local_client: bool,
|
||||
addresses: EngineZmqAddresses,
|
||||
executor_class: type[Any],
|
||||
log_stats: bool,
|
||||
dp_rank: int = 0,
|
||||
local_dp_rank: int = 0,
|
||||
):
|
||||
# Exercise the production Ray actor mixin without loading a model.
|
||||
EngineCoreActorMixin.__init__(
|
||||
self, vllm_config, addresses, dp_rank, local_dp_rank
|
||||
)
|
||||
|
||||
def _set_visible_devices(self, vllm_config: Any, local_dp_rank: int) -> None:
|
||||
pass
|
||||
|
||||
def wait_for_init(self) -> None:
|
||||
pass
|
||||
|
||||
def run(self) -> None:
|
||||
pass
|
||||
|
||||
def get_nixl_side_channel_host(self) -> str | None:
|
||||
return os.environ.get("VLLM_NIXL_SIDE_CHANNEL_HOST")
|
||||
|
||||
|
||||
class _DummyExecutor:
|
||||
pass
|
||||
|
||||
|
||||
def _make_vllm_config() -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
parallel_config=SimpleNamespace(
|
||||
data_parallel_size=1,
|
||||
data_parallel_size_local=1,
|
||||
enable_elastic_ep=False,
|
||||
world_size=1,
|
||||
),
|
||||
model_config=SimpleNamespace(is_moe=False),
|
||||
kv_transfer_config=None,
|
||||
)
|
||||
|
||||
|
||||
def _make_addresses() -> EngineZmqAddresses:
|
||||
return EngineZmqAddresses(
|
||||
inputs=["tcp://127.0.0.1:12345"],
|
||||
outputs=["tcp://127.0.0.1:12346"],
|
||||
)
|
||||
|
||||
|
||||
def _make_cpu_placement_group():
|
||||
pg = ray.util.placement_group(
|
||||
[{"CPU": 0.001}, {"CPU": 1.0}],
|
||||
strategy="PACK",
|
||||
)
|
||||
ray.get(pg.ready())
|
||||
return pg
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_context():
|
||||
started_ray = False
|
||||
if not ray.is_initialized():
|
||||
project_root = str(Path(__file__).resolve().parents[3])
|
||||
ray.init(
|
||||
num_cpus=2,
|
||||
runtime_env={"env_vars": {"PYTHONPATH": project_root}},
|
||||
log_to_driver=False,
|
||||
)
|
||||
started_ray = True
|
||||
|
||||
yield
|
||||
|
||||
if started_ray:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("ray_context")
|
||||
def test_driver_nixl_side_channel_host_does_not_leak_to_engine_core_actor(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
driver_marker = f"driver-only-nixl-host-{uuid.uuid4()}"
|
||||
created_placement_groups: list[Any] = []
|
||||
manager: CoreEngineActorManager | None = None
|
||||
|
||||
def create_dp_placement_groups(vllm_config: Any):
|
||||
pg = _make_cpu_placement_group()
|
||||
created_placement_groups.append(pg)
|
||||
return [pg], [0]
|
||||
|
||||
monkeypatch.setenv("VLLM_NIXL_SIDE_CHANNEL_HOST", driver_marker)
|
||||
monkeypatch.setattr("vllm.v1.engine.core.EngineCoreActor", _StubEngineCoreActor)
|
||||
monkeypatch.setattr(
|
||||
CoreEngineActorManager,
|
||||
"create_dp_placement_groups",
|
||||
staticmethod(create_dp_placement_groups),
|
||||
)
|
||||
|
||||
try:
|
||||
manager = CoreEngineActorManager(
|
||||
vllm_config=_make_vllm_config(),
|
||||
addresses=_make_addresses(),
|
||||
executor_class=_DummyExecutor,
|
||||
log_stats=False,
|
||||
)
|
||||
actor = manager.local_engine_actors[0]
|
||||
actor_host = ray.get(actor.get_nixl_side_channel_host.remote())
|
||||
node_host = ray.util.get_node_ip_address()
|
||||
|
||||
assert actor_host != driver_marker
|
||||
assert actor_host == node_host
|
||||
finally:
|
||||
if manager is not None:
|
||||
manager.shutdown()
|
||||
else:
|
||||
for pg in created_placement_groups:
|
||||
ray.util.remove_placement_group(pg)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user