forked from Karylab-cklius/vllm
Compare commits
57
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
42b812edf5 | ||
|
|
0a62f5eec9 | ||
|
|
3b1ef03be4 | ||
|
|
3c413a5481 | ||
|
|
79fd1bc7ed | ||
|
|
cee6751e54 | ||
|
|
16863072ca | ||
|
|
d628a3c5cb | ||
|
|
74dffae666 | ||
|
|
97c4317bf5 | ||
|
|
f6e868fbdf | ||
|
|
13bf242100 | ||
|
|
140dc2ec30 | ||
|
|
9ce74042d3 | ||
|
|
a8c13d2837 | ||
|
|
446bbb6ad3 | ||
|
|
2e6f2b451e | ||
|
|
92def124bc | ||
|
|
85b2fecab7 | ||
|
|
503697c9ce | ||
|
|
71bcd02ef3 | ||
|
|
dcacdf9a88 | ||
|
|
18f6bf5a21 | ||
|
|
07534b8782 | ||
|
|
3d635c58c0 | ||
|
+4 |
ebeb09d822 | ||
|
|
184577ae46 | ||
|
|
8c4fc4202a | ||
|
|
fe8b42e80c | ||
|
|
fe5b4e0fe7 | ||
|
|
0ce6613b9c | ||
|
|
379f0ec369 | ||
|
|
67c89fe40a | ||
|
|
d9b4990783 | ||
|
|
4d591db470 | ||
|
|
6ff7405b81 | ||
|
|
bcb9c133ba | ||
|
|
c8a6e272e0 | ||
|
|
a1b2d87498 | ||
|
|
418ba8ef14 | ||
|
|
5a6a9fc6f6 | ||
|
|
289cee0473 | ||
|
|
6ccb10d794 | ||
|
|
7a9cc5e7f0 | ||
|
|
d077622d60 | ||
|
|
dd6b3a5ef5 | ||
|
|
593d5a4033 | ||
|
|
6427603ae8 | ||
|
|
206eaed08d | ||
|
|
8f89381fc6 | ||
|
|
a7b801e26d | ||
|
|
4df1be9547 | ||
|
|
bc03f280c8 | ||
|
|
997132911e | ||
|
|
fc8bf6eedb | ||
|
|
07a40ede19 | ||
|
|
e1c8776e90 |
@@ -8,6 +8,7 @@ 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,6 +28,7 @@ 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"
|
||||
|
||||
@@ -41,6 +42,7 @@ 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"
|
||||
|
||||
@@ -54,6 +56,7 @@ 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"
|
||||
|
||||
@@ -67,6 +70,7 @@ 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"
|
||||
|
||||
@@ -80,6 +84,7 @@ 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"
|
||||
|
||||
@@ -93,6 +98,7 @@ 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"
|
||||
|
||||
@@ -138,6 +144,7 @@ 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: ~
|
||||
@@ -160,6 +167,7 @@ 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: ~
|
||||
@@ -184,6 +192,7 @@ 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: ~
|
||||
@@ -205,6 +214,7 @@ 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: ~
|
||||
@@ -231,6 +241,7 @@ 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: ~
|
||||
@@ -255,6 +266,7 @@ 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: ~
|
||||
@@ -280,6 +292,7 @@ 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: ~
|
||||
@@ -303,6 +316,7 @@ 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
|
||||
@@ -320,6 +334,7 @@ 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"
|
||||
|
||||
@@ -339,6 +354,7 @@ 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"
|
||||
|
||||
@@ -356,15 +372,7 @@ 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"
|
||||
|
||||
- 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"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "Manifest: CUDA 13.0" "public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"'
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 12.9"
|
||||
depends_on:
|
||||
@@ -377,6 +385,7 @@ 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:
|
||||
@@ -389,6 +398,7 @@ 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:
|
||||
@@ -401,6 +411,7 @@ 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:
|
||||
@@ -438,59 +449,6 @@ 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)
|
||||
# =============================================================================
|
||||
@@ -604,7 +562,7 @@ steps:
|
||||
echo ""
|
||||
echo " Build complete - Image and wheels cached"
|
||||
fi
|
||||
|
||||
|
||||
artifact_paths:
|
||||
- "artifacts/rocm-base-wheels/*.whl"
|
||||
env:
|
||||
@@ -820,7 +778,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"
|
||||
@@ -847,3 +805,60 @@ 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"
|
||||
|
||||
+9
@@ -0,0 +1,9 @@
|
||||
#!/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'
|
||||
@@ -1,27 +0,0 @@
|
||||
#!/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,11 +39,17 @@ fi
|
||||
|
||||
set -x # avoid printing secrets above
|
||||
|
||||
# install twine and sdist build prerequisites from pypi
|
||||
python3 -m venv /tmp/vllm-release-env
|
||||
# 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
|
||||
source /tmp/vllm-release-env/bin/activate
|
||||
pip install twine
|
||||
pip install -r requirements/build/cuda.txt
|
||||
uv pip install twine
|
||||
uv pip install -r requirements/build/cuda.txt
|
||||
python3 -m twine --version
|
||||
|
||||
# copy release wheels to local directory
|
||||
|
||||
@@ -91,6 +91,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,7 +106,11 @@ 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
|
||||
@@ -115,6 +120,7 @@ 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
|
||||
|
||||
@@ -114,6 +114,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
|
||||
|
||||
@@ -79,6 +79,7 @@ steps:
|
||||
- tests/v1/core
|
||||
- tests/v1/executor
|
||||
- tests/v1/kv_offload
|
||||
- tests/v1/simple_kv_offload
|
||||
- tests/v1/worker
|
||||
- tests/v1/kv_connector/unit
|
||||
- tests/v1/metrics
|
||||
@@ -90,6 +91,7 @@ steps:
|
||||
- pytest -v -s -m 'not cpu_test' v1/core
|
||||
- pytest -v -s v1/executor
|
||||
- pytest -v -s v1/kv_offload
|
||||
- pytest -v -s v1/simple_kv_offload
|
||||
- pytest -v -s v1/worker
|
||||
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/metrics
|
||||
@@ -212,6 +214,7 @@ steps:
|
||||
depends_on: ~
|
||||
optional: true
|
||||
timeout_in_minutes: 20
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- tests/standalone_tests/python_only_compile.sh
|
||||
- setup.py
|
||||
|
||||
@@ -106,13 +106,12 @@ steps:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- tests/v1/spec_decode/test_max_len.py
|
||||
- tests/v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
|
||||
- tests/v1/spec_decode/test_synthetic_rejection_sampler_utils.py
|
||||
- tests/v1/spec_decode/test_rejection_sampler_utils.py
|
||||
- tests/v1/e2e/spec_decode/test_spec_decode.py
|
||||
commands:
|
||||
- set -x
|
||||
- export VLLM_USE_V2_MODEL_RUNNER=1
|
||||
- pytest -v -s v1/spec_decode/test_max_len.py -k "eagle or mtp"
|
||||
- pytest -v -s v1/spec_decode/test_probabilistic_rejection_sampler_utils.py
|
||||
- pytest -v -s v1/spec_decode/test_rejection_sampler_utils.py
|
||||
- pytest -v -s v1/spec_decode/test_synthetic_rejection_sampler_utils.py
|
||||
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle or mtp"
|
||||
|
||||
@@ -5,6 +5,7 @@ steps:
|
||||
- label: Basic Models Tests (Initialization)
|
||||
key: basic-models-tests-initialization
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
torch_nightly: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -84,19 +84,19 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
|
||||
|
||||
- label: DFlash Speculators Correctness
|
||||
key: dflash-speculators-correctness
|
||||
timeout_in_minutes: 30
|
||||
- label: Speculators Correctness
|
||||
key: speculators-correctness
|
||||
timeout_in_minutes: 60
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/model_executor/models/qwen3_dflash.py
|
||||
- tests/v1/spec_decode/test_speculators_dflash.py
|
||||
- tests/v1/spec_decode/test_speculators_correctness.py
|
||||
commands:
|
||||
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
|
||||
- pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
|
||||
- pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
|
||||
|
||||
- label: Spec Decode MTP hybrid (B200)
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -9,6 +9,9 @@ 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
|
||||
|
||||
+71
-70
@@ -369,7 +369,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
FetchContent_MakeAvailable(cutlass)
|
||||
|
||||
list(APPEND VLLM_EXT_SRC
|
||||
"csrc/quantization/awq/gemm_kernels.cu"
|
||||
"csrc/cutlass_extensions/common.cpp")
|
||||
|
||||
set_gencode_flags_for_srcs(
|
||||
@@ -501,46 +500,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
" in CUDA target architectures")
|
||||
endif()
|
||||
|
||||
# Only build AllSpark kernels if we are building for at least some compatible archs.
|
||||
cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
|
||||
if (ALLSPARK_ARCHS)
|
||||
set(ALLSPARK_SRCS
|
||||
"csrc/quantization/gptq_allspark/allspark_repack.cu"
|
||||
"csrc/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${ALLSPARK_SRCS}"
|
||||
CUDA_ARCHS "${ALLSPARK_ARCHS}")
|
||||
list(APPEND VLLM_EXT_SRC "${ALLSPARK_SRCS}")
|
||||
message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building AllSpark kernels as no compatible archs found"
|
||||
" in CUDA target architectures")
|
||||
endif()
|
||||
|
||||
# CUTLASS MLA Archs and flags
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
|
||||
set(SRCS
|
||||
"csrc/attention/mla/sm100_cutlass_mla_kernel.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SRCS}"
|
||||
CUDA_ARCHS "${MLA_ARCHS}")
|
||||
list(APPEND VLLM_EXT_SRC "${SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
|
||||
# Add MLA-specific include directories only to MLA source files
|
||||
set_source_files_properties(${SRCS}
|
||||
PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
|
||||
message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building CUTLASS MLA as no compatible archs were found.")
|
||||
# clear MLA_ARCHS
|
||||
set(MLA_ARCHS)
|
||||
endif()
|
||||
|
||||
# Expert-specialization MXFP8 blockscaled grouped kernels (SM100+).
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(ES_MXFP8_GROUPED_MM_ARCHS "10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
@@ -568,24 +527,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
|
||||
set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${DSV3_FUSED_A_GEMM_SRC}"
|
||||
CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
|
||||
list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
|
||||
message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
|
||||
"in CUDA target architectures.")
|
||||
endif()
|
||||
|
||||
#
|
||||
# Machete kernels
|
||||
|
||||
@@ -657,16 +598,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
|
||||
|
||||
# Hadacore kernels
|
||||
cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
|
||||
if(HADACORE_ARCHS)
|
||||
set(SRCS "csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SRCS}"
|
||||
CUDA_ARCHS "${HADACORE_ARCHS}")
|
||||
list(APPEND VLLM_EXT_SRC "${SRCS}")
|
||||
message(STATUS "Building hadacore")
|
||||
endif()
|
||||
|
||||
# if CUDA endif
|
||||
endif()
|
||||
@@ -716,7 +647,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC
|
||||
"csrc/libtorch_stable/permute_cols.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu")
|
||||
"csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"
|
||||
"csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")
|
||||
endif()
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
@@ -725,6 +657,40 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
CUDA_ARCHS "${CUDA_ARCHS}")
|
||||
endif()
|
||||
|
||||
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
|
||||
set(SRCS "csrc/libtorch_stable/dsv3_fused_a_gemm.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SRCS}"
|
||||
CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
|
||||
message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
|
||||
"in CUDA target architectures.")
|
||||
endif()
|
||||
|
||||
# Only build AllSpark kernels if we are building for at least some compatible archs.
|
||||
cuda_archs_loose_intersection(ALLSPARK_ARCHS "8.0;8.6;8.7;8.9" "${CUDA_ARCHS}")
|
||||
if (ALLSPARK_ARCHS)
|
||||
set(SRCS
|
||||
"csrc/libtorch_stable/quantization/gptq_allspark/allspark_repack.cu"
|
||||
"csrc/libtorch_stable/quantization/gptq_allspark/allspark_qgemm_w8a16.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SRCS}"
|
||||
CUDA_ARCHS "${ALLSPARK_ARCHS}")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
|
||||
message(STATUS "Building AllSpark kernels for archs: ${ALLSPARK_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building AllSpark kernels as no compatible archs found"
|
||||
" in CUDA target architectures")
|
||||
endif()
|
||||
|
||||
#
|
||||
# CUTLASS scaled_mm kernels (moved from _C to _C_stable_libtorch)
|
||||
#
|
||||
@@ -1034,6 +1000,41 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# CUTLASS MLA Archs and flags
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(MLA_ARCHS "10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(MLA_ARCHS "10.0a;10.1a;10.3a;12.0a;12.1a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND MLA_ARCHS)
|
||||
set(SRCS
|
||||
"csrc/libtorch_stable/attention/mla/sm100_cutlass_mla_kernel.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SRCS}"
|
||||
CUDA_ARCHS "${MLA_ARCHS}")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MLA=1")
|
||||
# Add MLA-specific include directories only to MLA source files
|
||||
set_source_files_properties(${SRCS}
|
||||
PROPERTIES INCLUDE_DIRECTORIES "${CUTLASS_DIR}/examples/77_blackwell_fmha;${CUTLASS_DIR}/examples/common")
|
||||
message(STATUS "Building CUTLASS MLA for archs: ${MLA_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building CUTLASS MLA as no compatible archs were found.")
|
||||
# clear MLA_ARCHS
|
||||
set(MLA_ARCHS)
|
||||
endif()
|
||||
|
||||
# Hadacore kernels
|
||||
cuda_archs_loose_intersection(HADACORE_ARCHS "8.0+PTX;9.0+PTX" "${CUDA_ARCHS}")
|
||||
if(HADACORE_ARCHS)
|
||||
set(SRCS "csrc/libtorch_stable/quantization/hadamard/hadacore/hadamard_transform_cuda.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${SRCS}"
|
||||
CUDA_ARCHS "${HADACORE_ARCHS}")
|
||||
list(APPEND VLLM_STABLE_EXT_SRC "${SRCS}")
|
||||
message(STATUS "Building hadacore")
|
||||
endif()
|
||||
|
||||
message(STATUS "Enabling C_stable extension.")
|
||||
define_extension_target(
|
||||
_C_stable_libtorch
|
||||
|
||||
@@ -27,10 +27,10 @@ from vllm.model_executor.layers.fused_moe.config import (
|
||||
RoutingMethodType,
|
||||
_get_config_dtype_str,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import *
|
||||
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
|
||||
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
|
||||
TritonOrDeepGemmExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import *
|
||||
from vllm.transformers_utils.config import get_config
|
||||
from vllm.triton_utils import triton
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
@@ -53,48 +53,80 @@ cuda_archs_loose_intersection(DEEPGEMM_ARCHS
|
||||
if(DEEPGEMM_ARCHS)
|
||||
message(STATUS "DeepGEMM CUDA architectures: ${DEEPGEMM_ARCHS}")
|
||||
|
||||
find_package(CUDAToolkit REQUIRED)
|
||||
|
||||
#
|
||||
# Build the _C pybind11 extension from DeepGEMM's C++ source.
|
||||
# This is a CXX-only module — CUDA kernels are JIT-compiled at runtime.
|
||||
# DeepGEMM integration notes
|
||||
# --------------------------
|
||||
# We vendor DeepGEMM into vllm/third_party/deep_gemm/ and bundle a
|
||||
# `_C.cpython-X.Y-*.so` for every CPython in `requires-python`. The
|
||||
# per-Python build is delegated to tools/build_deepgemm_C.py.
|
||||
#
|
||||
# Why per-Python: DeepGEMM's binding uses PYBIND11_MODULE, which links
|
||||
# private CPython symbols — a single `_C.abi3.so` is not viable today
|
||||
# (see #41476 / #41512 for the failed attempt).
|
||||
#
|
||||
# TODOs (tracked in vllm-project/vllm#42431):
|
||||
# - Replace DeepGEMM's pybind11 binding with a TORCH_LIBRARY + shim
|
||||
# binding (cf. vllm-flash-attention/csrc/common/pytorch_shim.h) to
|
||||
# collapse to one `_C.abi3.so`. Needs either an upstream change or
|
||||
# a maintained binding fork in vLLM.
|
||||
# - AOT-compile DeepGEMM's CUDA kernels instead of runtime JIT to drop
|
||||
# the vendored CUTLASS/CCCL headers and the CUDA-toolkit-at-runtime
|
||||
# requirement.
|
||||
#
|
||||
Python_add_library(_deep_gemm_C MODULE WITH_SOABI
|
||||
"${deepgemm_SOURCE_DIR}/csrc/python_api.cpp")
|
||||
|
||||
# The pybind11 module name must be _C to match DeepGEMM's Python imports.
|
||||
set_target_properties(_deep_gemm_C PROPERTIES OUTPUT_NAME "_C")
|
||||
# DEEPGEMM_PYTHON_INTERPRETERS: ":"-separated target Python paths.
|
||||
# Empty/unset → fall back to the build interpreter (editable installs).
|
||||
# (Empty-but-set env vars test as DEFINED in cmake — treat as unset.)
|
||||
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}")
|
||||
|
||||
target_compile_definitions(_deep_gemm_C PRIVATE
|
||||
"-DTORCH_EXTENSION_NAME=_C")
|
||||
# add_custom_command does no implicit header scanning; glob explicitly so
|
||||
# header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild.
|
||||
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")
|
||||
|
||||
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)
|
||||
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 interpreters that resolve to the same CPython.
|
||||
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})
|
||||
|
||||
#
|
||||
# Vendor DeepGEMM Python package files
|
||||
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG f5bc33cfc02c744d24a2e9d50e6db656de40611c
|
||||
GIT_TAG bce29425653ec0fbc579d329883030e832d15ada
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
+26
-18
@@ -1,7 +1,13 @@
|
||||
#pragma once
|
||||
|
||||
// For TORCH_CHECK
|
||||
#include <torch/library.h>
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <tuple>
|
||||
#include <utility>
|
||||
#include <variant>
|
||||
|
||||
// For STD_TORCH_CHECK
|
||||
#include <torch/headeronly/util/Exception.h>
|
||||
|
||||
namespace vllm {
|
||||
|
||||
@@ -45,7 +51,7 @@ class ScalarType {
|
||||
// IEEE 754 compliant floating point type
|
||||
static constexpr ScalarType float_IEEE754(uint8_t exponent,
|
||||
uint8_t mantissa) {
|
||||
TORCH_CHECK(mantissa > 0 && exponent > 0);
|
||||
STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
|
||||
return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754);
|
||||
}
|
||||
|
||||
@@ -53,11 +59,12 @@ class ScalarType {
|
||||
static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa,
|
||||
bool finite_values_only,
|
||||
NanRepr nan_repr) {
|
||||
TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
|
||||
TORCH_CHECK(mantissa > 0 && exponent > 0);
|
||||
TORCH_CHECK(nan_repr != NAN_IEEE_754,
|
||||
"use `float_IEEE754` constructor for floating point types that "
|
||||
"follow IEEE 754 conventions");
|
||||
STD_TORCH_CHECK(nan_repr < NAN_REPR_ID_MAX, "Invalid NanRepr");
|
||||
STD_TORCH_CHECK(mantissa > 0 && exponent > 0);
|
||||
STD_TORCH_CHECK(
|
||||
nan_repr != NAN_IEEE_754,
|
||||
"use `float_IEEE754` constructor for floating point types that "
|
||||
"follow IEEE 754 conventions");
|
||||
return ScalarType(exponent, mantissa, true, 0, finite_values_only,
|
||||
nan_repr);
|
||||
}
|
||||
@@ -176,8 +183,8 @@ class ScalarType {
|
||||
|
||||
private:
|
||||
double _floating_point_max() const {
|
||||
TORCH_CHECK(mantissa <= 52 && exponent <= 11,
|
||||
"Cannot represent max/min as a double for type ", str());
|
||||
STD_TORCH_CHECK(mantissa <= 52 && exponent <= 11,
|
||||
"Cannot represent max/min as a double for type ", str());
|
||||
|
||||
uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1;
|
||||
if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) {
|
||||
@@ -186,8 +193,8 @@ class ScalarType {
|
||||
|
||||
uint64_t max_exponent = (uint64_t(1) << exponent) - 2;
|
||||
if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) {
|
||||
TORCH_CHECK(exponent < 11,
|
||||
"Cannot represent max/min as a double for type ", str());
|
||||
STD_TORCH_CHECK(exponent < 11,
|
||||
"Cannot represent max/min as a double for type ", str());
|
||||
max_exponent += 1;
|
||||
}
|
||||
|
||||
@@ -216,16 +223,17 @@ class ScalarType {
|
||||
if (is_floating_point()) {
|
||||
return {_floating_point_max()};
|
||||
} else {
|
||||
TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
|
||||
"Cannot represent max as a int64_t");
|
||||
STD_TORCH_CHECK(size_bits() < 64 || size_bits() == 64 && is_signed(),
|
||||
"Cannot represent max as a int64_t");
|
||||
return {(int64_t(1) << mantissa) - 1};
|
||||
}
|
||||
}
|
||||
|
||||
constexpr std::variant<int64_t, double> _raw_min() const {
|
||||
if (is_floating_point()) {
|
||||
TORCH_CHECK(is_signed(),
|
||||
"We currently assume all floating point types are signed");
|
||||
STD_TORCH_CHECK(
|
||||
is_signed(),
|
||||
"We currently assume all floating point types are signed");
|
||||
constexpr uint64_t sign_bit_double = (uint64_t(1) << 63);
|
||||
|
||||
double max = _floating_point_max();
|
||||
@@ -233,8 +241,8 @@ class ScalarType {
|
||||
uint64_t min_raw = max_raw | sign_bit_double;
|
||||
return {*reinterpret_cast<double*>(&min_raw)};
|
||||
} else {
|
||||
TORCH_CHECK(!is_signed() || size_bits() <= 64,
|
||||
"Cannot represent min as a int64_t");
|
||||
STD_TORCH_CHECK(!is_signed() || size_bits() <= 64,
|
||||
"Cannot represent min as a int64_t");
|
||||
if (is_signed()) {
|
||||
// set the top bit to 1 (i.e. INT64_MIN) and the rest to 0
|
||||
// then perform an arithmetic shift right to set all the bits above
|
||||
|
||||
+55
-56
@@ -18,13 +18,12 @@ limitations under the License.
|
||||
* Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
|
||||
* by Alcanderian JieXin Liang
|
||||
*/
|
||||
#include "core/registration.h"
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include <torch/csrc/stable/library.h>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <cutlass/cutlass.h>
|
||||
#include <cutlass/kernel_hardware_info.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <cute/tensor.hpp>
|
||||
#include <iostream>
|
||||
@@ -35,27 +34,27 @@ limitations under the License.
|
||||
// clang-format off
|
||||
#if !defined(CUDA_VERSION) || CUDA_VERSION < 12040
|
||||
void sm100_cutlass_mla_decode(
|
||||
torch::Tensor const& out,
|
||||
torch::Tensor const& lse,
|
||||
torch::Tensor const& q_nope,
|
||||
torch::Tensor const& q_pe,
|
||||
torch::Tensor const& kv_c_and_k_pe_cache,
|
||||
torch::Tensor const& seq_lens,
|
||||
torch::Tensor const& page_table,
|
||||
torch::Tensor const& workspace,
|
||||
torch::stable::Tensor const& out,
|
||||
torch::stable::Tensor const& lse,
|
||||
torch::stable::Tensor const& q_nope,
|
||||
torch::stable::Tensor const& q_pe,
|
||||
torch::stable::Tensor const& kv_c_and_k_pe_cache,
|
||||
torch::stable::Tensor const& seq_lens,
|
||||
torch::stable::Tensor const& page_table,
|
||||
torch::stable::Tensor const& workspace,
|
||||
double sm_scale,
|
||||
int64_t num_kv_splits) {
|
||||
TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
|
||||
STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_decode");
|
||||
}
|
||||
int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_batches, int64_t sm_count, int64_t num_kv_splits) {
|
||||
TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
|
||||
STD_TORCH_CHECK(false, "CUDA version must be >= 12.4 for cutlass_mla_get_workspace_size");
|
||||
}
|
||||
#else
|
||||
|
||||
#define CUTLASS_CHECK(status) \
|
||||
{ \
|
||||
cutlass::Status error = status; \
|
||||
TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
|
||||
STD_TORCH_CHECK(error == cutlass::Status::kSuccess, cutlassGetStatusString(error)); \
|
||||
}
|
||||
|
||||
using namespace cute;
|
||||
@@ -100,23 +99,23 @@ struct MlaSm100 {
|
||||
|
||||
template <typename T>
|
||||
typename T::Fmha::Arguments args_from_options(
|
||||
at::Tensor const& out,
|
||||
at::Tensor const& lse,
|
||||
at::Tensor const& q_nope,
|
||||
at::Tensor const& q_pe,
|
||||
at::Tensor const& kv_c_and_k_pe_cache,
|
||||
at::Tensor const& seq_lens,
|
||||
at::Tensor const& page_table,
|
||||
torch::stable::Tensor const& out,
|
||||
torch::stable::Tensor const& lse,
|
||||
torch::stable::Tensor const& q_nope,
|
||||
torch::stable::Tensor const& q_pe,
|
||||
torch::stable::Tensor const& kv_c_and_k_pe_cache,
|
||||
torch::stable::Tensor const& seq_lens,
|
||||
torch::stable::Tensor const& page_table,
|
||||
double sm_scale,
|
||||
int64_t num_kv_splits) {
|
||||
cutlass::KernelHardwareInfo hw_info;
|
||||
hw_info.device_id = q_nope.device().index();
|
||||
hw_info.device_id = q_nope.get_device_index();
|
||||
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
|
||||
|
||||
int batches = q_nope.sizes()[0];
|
||||
int page_count_per_seq = page_table.sizes()[1];
|
||||
int page_count_total = kv_c_and_k_pe_cache.sizes()[0];
|
||||
int page_size = kv_c_and_k_pe_cache.sizes()[1];
|
||||
int batches = q_nope.size(0);
|
||||
int page_count_per_seq = page_table.size(1);
|
||||
int page_count_total = kv_c_and_k_pe_cache.size(0);
|
||||
int page_size = kv_c_and_k_pe_cache.size(1);
|
||||
int max_seq_len = page_size * page_count_per_seq;
|
||||
using TileShapeH = typename T::TileShapeH;
|
||||
using TileShapeD = typename T::TileShapeD;
|
||||
@@ -186,14 +185,14 @@ typename T::Fmha::Arguments args_from_options(
|
||||
|
||||
template <typename Element, typename ElementOut, bool IsPaged128, typename PersistenceOption>
|
||||
void runMla(
|
||||
at::Tensor const& out,
|
||||
at::Tensor const& lse,
|
||||
at::Tensor const& q_nope,
|
||||
at::Tensor const& q_pe,
|
||||
at::Tensor const& kv_c_and_k_pe_cache,
|
||||
at::Tensor const& seq_lens,
|
||||
at::Tensor const& page_table,
|
||||
at::Tensor const& workspace,
|
||||
torch::stable::Tensor const& out,
|
||||
torch::stable::Tensor const& lse,
|
||||
torch::stable::Tensor const& q_nope,
|
||||
torch::stable::Tensor const& q_pe,
|
||||
torch::stable::Tensor const& kv_c_and_k_pe_cache,
|
||||
torch::stable::Tensor const& seq_lens,
|
||||
torch::stable::Tensor const& page_table,
|
||||
torch::stable::Tensor const& workspace,
|
||||
double sm_scale,
|
||||
int64_t num_kv_splits,
|
||||
cudaStream_t stream) {
|
||||
@@ -220,37 +219,37 @@ void runMla(
|
||||
}()
|
||||
|
||||
void sm100_cutlass_mla_decode(
|
||||
torch::Tensor const& out,
|
||||
torch::Tensor const& lse,
|
||||
torch::Tensor const& q_nope,
|
||||
torch::Tensor const& q_pe,
|
||||
torch::Tensor const& kv_c_and_k_pe_cache,
|
||||
torch::Tensor const& seq_lens,
|
||||
torch::Tensor const& page_table,
|
||||
torch::Tensor const& workspace,
|
||||
torch::stable::Tensor const& out,
|
||||
torch::stable::Tensor const& lse,
|
||||
torch::stable::Tensor const& q_nope,
|
||||
torch::stable::Tensor const& q_pe,
|
||||
torch::stable::Tensor const& kv_c_and_k_pe_cache,
|
||||
torch::stable::Tensor const& seq_lens,
|
||||
torch::stable::Tensor const& page_table,
|
||||
torch::stable::Tensor const& workspace,
|
||||
double sm_scale,
|
||||
int64_t num_kv_splits) {
|
||||
auto in_dtype = q_nope.dtype();
|
||||
at::cuda::CUDAGuard device_guard{(char)q_nope.get_device()};
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(q_nope.get_device());
|
||||
const int page_size = kv_c_and_k_pe_cache.sizes()[1];
|
||||
|
||||
auto in_dtype = q_nope.scalar_type();
|
||||
torch::stable::accelerator::DeviceGuard device_guard(q_nope.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream(q_nope.get_device_index());
|
||||
const int page_size = kv_c_and_k_pe_cache.size(1);
|
||||
|
||||
// NOTE(alcanderian): IsPersistent has bug with manual split_kv.
|
||||
// Kernel will hang if batch is too large with large num_kv_splits. (for example bs=8, num_kv_splits=8)
|
||||
// Maybe per batch split kv will fix this.
|
||||
DISPATCH_BOOL(page_size == 128, IsPaged128, [&] {
|
||||
DISPATCH_BOOL(num_kv_splits <= 1, NotManualSplitKV, [&] {
|
||||
if (in_dtype == at::ScalarType::Half) {
|
||||
if (in_dtype == torch::headeronly::ScalarType::Half) {
|
||||
runMla<cutlass::half_t, cutlass::half_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
|
||||
out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
|
||||
} else if (in_dtype == at::ScalarType::BFloat16) {
|
||||
} else if (in_dtype == torch::headeronly::ScalarType::BFloat16) {
|
||||
runMla<cutlass::bfloat16_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
|
||||
out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
|
||||
} else if (in_dtype == at::ScalarType::Float8_e4m3fn) {
|
||||
} else if (in_dtype == torch::headeronly::ScalarType::Float8_e4m3fn) {
|
||||
runMla<cutlass::float_e4m3_t, cutlass::bfloat16_t, IsPaged128, IsPersistent<NotManualSplitKV>>(
|
||||
out, lse, q_nope, q_pe, kv_c_and_k_pe_cache, seq_lens, page_table, workspace, sm_scale, num_kv_splits, stream);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unsupported input data type of MLA");
|
||||
STD_TORCH_CHECK(false, "Unsupported input data type of MLA");
|
||||
}
|
||||
return true;
|
||||
});
|
||||
@@ -280,12 +279,12 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
|
||||
|
||||
#endif
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("sm100_cutlass_mla_decode", &sm100_cutlass_mla_decode);
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
|
||||
m.impl("sm100_cutlass_mla_decode", TORCH_BOX(&sm100_cutlass_mla_decode));
|
||||
}
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CatchAll, m) {
|
||||
m.impl("sm100_cutlass_mla_get_workspace_size", &sm100_cutlass_mla_get_workspace_size);
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CompositeExplicitAutograd, m) {
|
||||
m.impl("sm100_cutlass_mla_get_workspace_size", TORCH_BOX(&sm100_cutlass_mla_get_workspace_size));
|
||||
}
|
||||
|
||||
// clang-format on
|
||||
@@ -20,13 +20,15 @@
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/all.h>
|
||||
#include <torch/csrc/stable/library.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
|
||||
#include "core/registration.h"
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <cstdlib>
|
||||
#include <mutex>
|
||||
@@ -34,7 +36,7 @@
|
||||
namespace {
|
||||
|
||||
inline int getSMVersion() {
|
||||
auto* props = at::cuda::getCurrentDeviceProperties();
|
||||
auto* props = get_device_prop();
|
||||
return props->major * 10 + props->minor;
|
||||
}
|
||||
|
||||
@@ -700,37 +702,40 @@ template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
|
||||
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
|
||||
cudaStream_t);
|
||||
|
||||
void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
|
||||
torch::Tensor const& mat_b) {
|
||||
TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
|
||||
void dsv3_fused_a_gemm(torch::stable::Tensor& output,
|
||||
torch::stable::Tensor const& mat_a,
|
||||
torch::stable::Tensor const& mat_b) {
|
||||
STD_TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
|
||||
int const num_tokens = mat_a.size(0);
|
||||
int const hd_in = mat_a.size(1);
|
||||
int const hd_out = mat_b.size(1);
|
||||
|
||||
constexpr int kHdIn = 7168;
|
||||
constexpr int kHdOut = 2112;
|
||||
TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
|
||||
"required 1 <= mat_a.shape[0] <= 16")
|
||||
TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168")
|
||||
TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112")
|
||||
TORCH_CHECK(output.size(0) == num_tokens,
|
||||
"required output.shape[0] == mat_a.shape[0]")
|
||||
TORCH_CHECK(output.size(1) == hd_out,
|
||||
"required output.shape[1] == mat_b.shape[1]")
|
||||
STD_TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
|
||||
"required 1 <= mat_a.shape[0] <= 16");
|
||||
STD_TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168");
|
||||
STD_TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112");
|
||||
STD_TORCH_CHECK(output.size(0) == num_tokens,
|
||||
"required output.shape[0] == mat_a.shape[0]");
|
||||
STD_TORCH_CHECK(output.size(1) == hd_out,
|
||||
"required output.shape[1] == mat_b.shape[1]");
|
||||
|
||||
TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
|
||||
TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
|
||||
TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
|
||||
STD_TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
|
||||
STD_TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
|
||||
STD_TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
|
||||
|
||||
TORCH_CHECK(mat_a.scalar_type() == torch::kBFloat16 &&
|
||||
mat_b.scalar_type() == torch::kBFloat16,
|
||||
"Only BFloat16 input dtype is supported")
|
||||
TORCH_CHECK(output.scalar_type() == torch::kBFloat16,
|
||||
"Only BFloat16 output dtype is supported")
|
||||
STD_TORCH_CHECK(
|
||||
mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16 &&
|
||||
mat_b.scalar_type() == torch::headeronly::ScalarType::BFloat16,
|
||||
"Only BFloat16 input dtype is supported");
|
||||
STD_TORCH_CHECK(
|
||||
output.scalar_type() == torch::headeronly::ScalarType::BFloat16,
|
||||
"Only BFloat16 output dtype is supported");
|
||||
|
||||
TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
|
||||
STD_TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
|
||||
|
||||
auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device());
|
||||
auto stream = get_current_cuda_stream(mat_a.get_device_index());
|
||||
if (num_tokens <= 8) {
|
||||
invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
|
||||
reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
|
||||
@@ -746,6 +751,6 @@ void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
|
||||
}
|
||||
}
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("dsv3_fused_a_gemm", &dsv3_fused_a_gemm);
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
|
||||
m.impl("dsv3_fused_a_gemm", TORCH_BOX(&dsv3_fused_a_gemm));
|
||||
}
|
||||
@@ -143,4 +143,26 @@ void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
|
||||
const torch::stable::Tensor& expert_offsets,
|
||||
const torch::stable::Tensor& sf_offsets);
|
||||
|
||||
// AWQ ops
|
||||
torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats,
|
||||
torch::stable::Tensor _kernel,
|
||||
torch::stable::Tensor _scaling_factors,
|
||||
torch::stable::Tensor _zeros,
|
||||
int64_t split_k_iters);
|
||||
|
||||
torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
|
||||
torch::stable::Tensor _scaling_factors,
|
||||
torch::stable::Tensor _zeros,
|
||||
int64_t split_k_iters, int64_t thx,
|
||||
int64_t thy);
|
||||
|
||||
// DSV3 fused A GEMM: conditionally compiled so declaration and impl
|
||||
// registration are in the source file (dsv3_fused_a_gemm.cu)
|
||||
|
||||
// AllSpark ops: declarations are in the source files
|
||||
// (allspark_repack.cu and allspark_qgemm_w8a16.cu)
|
||||
|
||||
#endif
|
||||
|
||||
torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
|
||||
bool inplace);
|
||||
|
||||
+41
-35
@@ -7,10 +7,11 @@ Shang and Dang, Xingyu and Han, Song}, journal={arXiv}, year={2023}
|
||||
}
|
||||
*/
|
||||
|
||||
#include <torch/all.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/csrc/stable/ops.h>
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
#include "dequantize.cuh"
|
||||
#include "libtorch_stable/quantization/awq/dequantize.cuh"
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
@@ -410,10 +411,11 @@ __global__ void __launch_bounds__(64)
|
||||
} // namespace awq
|
||||
} // namespace vllm
|
||||
|
||||
torch::Tensor awq_dequantize(torch::Tensor _kernel,
|
||||
torch::Tensor _scaling_factors,
|
||||
torch::Tensor _zeros, int64_t split_k_iters,
|
||||
int64_t thx, int64_t thy) {
|
||||
torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
|
||||
torch::stable::Tensor _scaling_factors,
|
||||
torch::stable::Tensor _zeros,
|
||||
int64_t split_k_iters, int64_t thx,
|
||||
int64_t thy) {
|
||||
int in_c = _kernel.size(0);
|
||||
int qout_c = _kernel.size(1);
|
||||
int out_c = qout_c * 8;
|
||||
@@ -437,23 +439,24 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
|
||||
y_blocks = (int)(in_c / 8);
|
||||
}
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(_scaling_factors));
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
_scaling_factors.get_device_index());
|
||||
|
||||
auto options = torch::TensorOptions()
|
||||
.dtype(_scaling_factors.dtype())
|
||||
.device(_scaling_factors.device());
|
||||
at::Tensor _de_kernel = torch::empty({in_c, out_c}, options);
|
||||
auto _de_kernel =
|
||||
torch::stable::empty({in_c, out_c}, _scaling_factors.scalar_type(),
|
||||
std::nullopt, _scaling_factors.device());
|
||||
|
||||
auto kernel = reinterpret_cast<int*>(_kernel.data_ptr<int>());
|
||||
auto de_kernel = reinterpret_cast<half*>(_de_kernel.data_ptr<at::Half>());
|
||||
auto scaling_factors =
|
||||
reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
|
||||
auto zeros = reinterpret_cast<int*>(_zeros.data_ptr<int>());
|
||||
auto kernel = reinterpret_cast<int*>(_kernel.mutable_data_ptr<int>());
|
||||
auto de_kernel = reinterpret_cast<half*>(
|
||||
_de_kernel.mutable_data_ptr<torch::headeronly::Half>());
|
||||
auto scaling_factors = reinterpret_cast<half*>(
|
||||
_scaling_factors.mutable_data_ptr<torch::headeronly::Half>());
|
||||
auto zeros = reinterpret_cast<int*>(_zeros.mutable_data_ptr<int>());
|
||||
|
||||
dim3 num_blocks(x_blocks, y_blocks);
|
||||
dim3 threads_per_block(x_thread, y_thread);
|
||||
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
const cudaStream_t stream = get_current_cuda_stream();
|
||||
vllm::awq::dequantize_weights<<<num_blocks, threads_per_block, 0, stream>>>(
|
||||
kernel, scaling_factors, zeros, de_kernel, G);
|
||||
|
||||
@@ -466,27 +469,30 @@ torch::Tensor awq_dequantize(torch::Tensor _kernel,
|
||||
// zeros: IC // G, OC // 8 [int32] -> cast to IC // G, OC [uint4b]
|
||||
// assume that batch_size < 16 for now
|
||||
|
||||
torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
|
||||
torch::Tensor _scaling_factors, torch::Tensor _zeros,
|
||||
int64_t split_k_iters) {
|
||||
torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats,
|
||||
torch::stable::Tensor _kernel,
|
||||
torch::stable::Tensor _scaling_factors,
|
||||
torch::stable::Tensor _zeros,
|
||||
int64_t split_k_iters) {
|
||||
int num_in_feats = _in_feats.size(0);
|
||||
int num_in_channels = _in_feats.size(1);
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(_in_feats));
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
_in_feats.get_device_index());
|
||||
|
||||
auto options = torch::TensorOptions()
|
||||
.dtype(_in_feats.dtype())
|
||||
.device(_in_feats.device());
|
||||
at::Tensor _out_feats =
|
||||
torch::empty({split_k_iters, num_in_feats, _kernel.size(1) * 8}, options);
|
||||
auto _out_feats = torch::stable::empty(
|
||||
{split_k_iters, num_in_feats, _kernel.size(1) * 8},
|
||||
_in_feats.scalar_type(), std::nullopt, _in_feats.device());
|
||||
int num_out_feats = _out_feats.size(-2);
|
||||
int num_out_channels = _out_feats.size(-1);
|
||||
|
||||
auto in_feats = reinterpret_cast<half*>(_in_feats.data_ptr<at::Half>());
|
||||
auto kernel = reinterpret_cast<int*>(_kernel.data_ptr<int>());
|
||||
auto out_feats = reinterpret_cast<half*>(_out_feats.data_ptr<at::Half>());
|
||||
auto scaling_factors =
|
||||
reinterpret_cast<half*>(_scaling_factors.data_ptr<at::Half>());
|
||||
auto zeros = reinterpret_cast<int*>(_zeros.data_ptr<int>());
|
||||
auto in_feats = reinterpret_cast<half*>(
|
||||
_in_feats.mutable_data_ptr<torch::headeronly::Half>());
|
||||
auto kernel = reinterpret_cast<int*>(_kernel.mutable_data_ptr<int>());
|
||||
auto out_feats = reinterpret_cast<half*>(
|
||||
_out_feats.mutable_data_ptr<torch::headeronly::Half>());
|
||||
auto scaling_factors = reinterpret_cast<half*>(
|
||||
_scaling_factors.mutable_data_ptr<torch::headeronly::Half>());
|
||||
auto zeros = reinterpret_cast<int*>(_zeros.mutable_data_ptr<int>());
|
||||
int group_size = num_in_channels / _scaling_factors.size(0);
|
||||
|
||||
if (num_out_channels % 64 != 0)
|
||||
@@ -498,7 +504,7 @@ torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
|
||||
if (num_out_channels % group_size != 0)
|
||||
throw std::invalid_argument("OC is not multiple of Group size");
|
||||
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
const cudaStream_t stream = get_current_cuda_stream();
|
||||
if (num_out_channels % 128 == 0) {
|
||||
int j_factors1 = num_out_channels / 128 / 1;
|
||||
dim3 num_blocks((num_out_feats + 16 - 1) / 16 * j_factors1 * split_k_iters);
|
||||
@@ -522,5 +528,5 @@ torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
|
||||
group_size, split_k_iters, in_feats, kernel, scaling_factors, zeros,
|
||||
num_in_feats, num_in_channels, num_out_channels, out_feats);
|
||||
}
|
||||
return _out_feats.sum(0);
|
||||
return torch::stable::sum(_out_feats, 0);
|
||||
}
|
||||
+56
-43
@@ -1,20 +1,28 @@
|
||||
#include "allspark_utils.cuh"
|
||||
#include <torch/all.h>
|
||||
#include "core/registration.h"
|
||||
|
||||
#include <torch/csrc/stable/c/shim.h>
|
||||
#include <torch/csrc/stable/library.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
|
||||
#include <cublas_v2.h>
|
||||
|
||||
at::Tensor as_g_workspace;
|
||||
#include "core/registration.h"
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
torch::stable::Tensor as_g_workspace;
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
|
||||
|
||||
torch::Tensor allspark_w8a16_gemm(
|
||||
torch::Tensor const& a, torch::Tensor const& b_qweight,
|
||||
torch::Tensor const& b_scales, std::optional<torch::Tensor> const& b_qzeros,
|
||||
int64_t n, int64_t group_size, int64_t sm_count, int64_t sm_version,
|
||||
torch::stable::Tensor allspark_w8a16_gemm(
|
||||
torch::stable::Tensor const& a, torch::stable::Tensor const& b_qweight,
|
||||
torch::stable::Tensor const& b_scales,
|
||||
std::optional<torch::stable::Tensor> const& b_qzeros, int64_t n,
|
||||
int64_t group_size, int64_t sm_count, int64_t sm_version,
|
||||
int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
STD_TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false, "allspark_w8a16_gemm(..) requires CUDA_ARCH >= 8.0");
|
||||
return torch::empty({1, 1});
|
||||
return torch::stable::empty({1, 1});
|
||||
}
|
||||
|
||||
#else
|
||||
@@ -848,8 +856,8 @@ void restore_N32_K16_dequantize_rhs_w8a16(const QT* qdata, const FT* scales,
|
||||
const int N_32align, const int N,
|
||||
const int K, const int GroupSize,
|
||||
cudaStream_t stream) {
|
||||
TORCH_CHECK(N % 8 == 0 && K % 16 == 0 && N_32align % 32 == 0,
|
||||
"Unsupported shape");
|
||||
STD_TORCH_CHECK(N % 8 == 0 && K % 16 == 0 && N_32align % 32 == 0,
|
||||
"Unsupported shape");
|
||||
if (GroupSize == -1) {
|
||||
const int BLOCK = 128;
|
||||
dim3 grid(N_32align / 32, ((K / 16) + 3) / 4);
|
||||
@@ -859,7 +867,7 @@ void restore_N32_K16_dequantize_rhs_w8a16(const QT* qdata, const FT* scales,
|
||||
}
|
||||
// TODO: Support SubChannel
|
||||
else {
|
||||
TORCH_CHECK(false, "Now only support PerChannel");
|
||||
STD_TORCH_CHECK(false, "Now only support PerChannel");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -916,24 +924,27 @@ void allspark_qgemm_w8a16_perc_ampere(
|
||||
|
||||
} // namespace allspark
|
||||
|
||||
torch::Tensor allspark_w8a16_gemm(
|
||||
torch::Tensor const& a, torch::Tensor const& b_qweight,
|
||||
torch::Tensor const& b_scales, std::optional<torch::Tensor> const& b_qzeros,
|
||||
int64_t n, int64_t group_size, int64_t sm_count, int64_t sm_version,
|
||||
torch::stable::Tensor allspark_w8a16_gemm(
|
||||
torch::stable::Tensor const& a, torch::stable::Tensor const& b_qweight,
|
||||
torch::stable::Tensor const& b_scales,
|
||||
std::optional<torch::stable::Tensor> const& b_qzeros, int64_t n,
|
||||
int64_t group_size, int64_t sm_count, int64_t sm_version,
|
||||
int64_t CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) {
|
||||
// Verify device and strides
|
||||
TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
|
||||
TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
|
||||
STD_TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
|
||||
STD_TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
|
||||
|
||||
TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
|
||||
TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
|
||||
STD_TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
|
||||
STD_TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
|
||||
|
||||
TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
|
||||
TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
|
||||
STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
|
||||
STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
|
||||
|
||||
if (has_zp) {
|
||||
TORCH_CHECK(b_qzeros.value().device().is_cuda(), "b_qzeros is not on GPU");
|
||||
TORCH_CHECK(b_qzeros.value().is_contiguous(), "b_qzeros is not contiguous");
|
||||
STD_TORCH_CHECK(b_qzeros.value().device().is_cuda(),
|
||||
"b_qzeros is not on GPU");
|
||||
STD_TORCH_CHECK(b_qzeros.value().is_contiguous(),
|
||||
"b_qzeros is not contiguous");
|
||||
}
|
||||
|
||||
int m = a.size(0);
|
||||
@@ -941,16 +952,17 @@ torch::Tensor allspark_w8a16_gemm(
|
||||
int k = a.size(1);
|
||||
|
||||
// Verify shape
|
||||
TORCH_CHECK(b_qweight.size(0) == n_32align,
|
||||
"Shape mismatch: b_qweight.size(0) = ", b_qweight.size(0),
|
||||
", n_32align = ", n_32align);
|
||||
TORCH_CHECK(b_qweight.size(1) == k,
|
||||
"Shape mismatch: b_qweight.size(1) = ", b_qweight.size(1),
|
||||
", k = ", k);
|
||||
STD_TORCH_CHECK(b_qweight.size(0) == n_32align,
|
||||
"Shape mismatch: b_qweight.size(0) = ", b_qweight.size(0),
|
||||
", n_32align = ", n_32align);
|
||||
STD_TORCH_CHECK(b_qweight.size(1) == k,
|
||||
"Shape mismatch: b_qweight.size(1) = ", b_qweight.size(1),
|
||||
", k = ", k);
|
||||
|
||||
TORCH_CHECK(group_size == -1, "Currently only supports group_size = -1");
|
||||
STD_TORCH_CHECK(group_size == -1, "Currently only supports group_size = -1");
|
||||
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(a));
|
||||
const torch::stable::accelerator::DeviceGuard device_guard(
|
||||
a.get_device_index());
|
||||
const void* a_ptr = reinterpret_cast<const void*>(a.data_ptr());
|
||||
const uint8_t* b_ptr = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
|
||||
const void* b_scale_ptr = reinterpret_cast<const void*>(b_scales.data_ptr());
|
||||
@@ -959,12 +971,12 @@ torch::Tensor allspark_w8a16_gemm(
|
||||
b_zero_ptr = reinterpret_cast<const void*>(b_qzeros.value().data_ptr());
|
||||
}
|
||||
|
||||
auto c_options = torch::TensorOptions().dtype(a.dtype()).device(a.device());
|
||||
torch::Tensor c = torch::empty({m, n}, c_options);
|
||||
void* c_ptr = reinterpret_cast<void*>(c.data_ptr());
|
||||
auto c =
|
||||
torch::stable::empty({m, n}, a.scalar_type(), std::nullopt, a.device());
|
||||
void* c_ptr = reinterpret_cast<void*>(c.mutable_data_ptr());
|
||||
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
|
||||
cudaStream_t stream = get_current_cuda_stream();
|
||||
cublasHandle_t handle = get_current_cuda_blas_handle();
|
||||
|
||||
allspark::BlockTileSplitkParams fused_gemm_params;
|
||||
|
||||
@@ -976,14 +988,15 @@ torch::Tensor allspark_w8a16_gemm(
|
||||
m, n, k, sm_count, fused_gemm_params);
|
||||
}
|
||||
|
||||
auto ws_options = torch::TensorOptions().dtype(at::kChar).device(a.device());
|
||||
if (as_g_workspace.numel() <
|
||||
ws_size) { // ws_options: kChar, so numel() is bytes
|
||||
as_g_workspace = torch::empty({long(ws_size)}, ws_options);
|
||||
as_g_workspace = torch::stable::empty({static_cast<int64_t>(ws_size)},
|
||||
torch::headeronly::ScalarType::Char,
|
||||
std::nullopt, a.device());
|
||||
}
|
||||
void* ws = reinterpret_cast<void*>(as_g_workspace.data_ptr());
|
||||
|
||||
if (a.dtype() == at::ScalarType::Half) {
|
||||
if (a.scalar_type() == torch::headeronly::ScalarType::Half) {
|
||||
allspark::allspark_qgemm_w8a16_perc_ampere<__half, uint8_t>(
|
||||
reinterpret_cast<const __half*>(a_ptr), b_ptr,
|
||||
reinterpret_cast<const __half*>(b_scale_ptr),
|
||||
@@ -991,7 +1004,7 @@ torch::Tensor allspark_w8a16_gemm(
|
||||
reinterpret_cast<__half*>(c_ptr), m, n_32align, n, k, ws,
|
||||
fused_gemm_params, group_size, CUBLAS_M_THRESHOLD, sm_version, stream,
|
||||
handle);
|
||||
} else if (a.dtype() == at::ScalarType::BFloat16) {
|
||||
} else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
allspark::allspark_qgemm_w8a16_perc_ampere<__nv_bfloat16, uint8_t>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(a_ptr), b_ptr,
|
||||
reinterpret_cast<const __nv_bfloat16*>(b_scale_ptr),
|
||||
@@ -1006,6 +1019,6 @@ torch::Tensor allspark_w8a16_gemm(
|
||||
|
||||
#endif
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("allspark_w8a16_gemm", &allspark_w8a16_gemm);
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
|
||||
m.impl("allspark_w8a16_gemm", TORCH_BOX(&allspark_w8a16_gemm));
|
||||
}
|
||||
+43
-32
@@ -1,6 +1,11 @@
|
||||
#include "allspark_utils.cuh"
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <torch/csrc/stable/library.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
|
||||
#include "core/registration.h"
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
|
||||
namespace allspark {
|
||||
|
||||
@@ -99,36 +104,40 @@ void rearrange_kn_weight_as_n32k16_order_ldg16(
|
||||
} // namespace allspark
|
||||
|
||||
void rearrange_kn_weight_as_n32k16_order(
|
||||
torch::Tensor const& b_qweight, torch::Tensor const& b_scales,
|
||||
std::optional<torch::Tensor> const& b_zeros, bool has_zp,
|
||||
torch::Tensor& b_qweight_reorder, torch::Tensor& b_scales_reorder,
|
||||
std::optional<torch::Tensor> const& b_zeros_reorder, const int64_t K,
|
||||
const int64_t N, const int64_t N_32align) {
|
||||
torch::stable::Tensor const& b_qweight,
|
||||
torch::stable::Tensor const& b_scales,
|
||||
std::optional<torch::stable::Tensor> const& b_zeros, bool has_zp,
|
||||
torch::stable::Tensor& b_qweight_reorder,
|
||||
torch::stable::Tensor& b_scales_reorder,
|
||||
std::optional<torch::stable::Tensor> const& b_zeros_reorder,
|
||||
const int64_t K, const int64_t N, const int64_t N_32align) {
|
||||
// Verify device and strides
|
||||
TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
|
||||
TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
|
||||
STD_TORCH_CHECK(b_qweight.device().is_cuda(), "b_qweight is not on GPU");
|
||||
STD_TORCH_CHECK(b_qweight.is_contiguous(), "b_qweight is not contiguous");
|
||||
|
||||
TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
|
||||
TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
|
||||
STD_TORCH_CHECK(b_scales.device().is_cuda(), "b_scales is not on GPU");
|
||||
STD_TORCH_CHECK(b_scales.is_contiguous(), "b_scales is not contiguous");
|
||||
|
||||
TORCH_CHECK(b_qweight_reorder.device().is_cuda(),
|
||||
"b_qweight_reorder is not on GPU");
|
||||
TORCH_CHECK(b_qweight_reorder.is_contiguous(),
|
||||
"b_qweight_reorder is not contiguous");
|
||||
STD_TORCH_CHECK(b_qweight_reorder.device().is_cuda(),
|
||||
"b_qweight_reorder is not on GPU");
|
||||
STD_TORCH_CHECK(b_qweight_reorder.is_contiguous(),
|
||||
"b_qweight_reorder is not contiguous");
|
||||
|
||||
TORCH_CHECK(b_scales_reorder.device().is_cuda(),
|
||||
"b_scales_reorder is not on GPU");
|
||||
TORCH_CHECK(b_scales_reorder.is_contiguous(),
|
||||
"b_scales_reorder is not contiguous");
|
||||
STD_TORCH_CHECK(b_scales_reorder.device().is_cuda(),
|
||||
"b_scales_reorder is not on GPU");
|
||||
STD_TORCH_CHECK(b_scales_reorder.is_contiguous(),
|
||||
"b_scales_reorder is not contiguous");
|
||||
|
||||
if (has_zp) {
|
||||
TORCH_CHECK(b_zeros.value().device().is_cuda(), "b_zeros is not on GPU");
|
||||
TORCH_CHECK(b_zeros.value().is_contiguous(), "b_zeros is not contiguous");
|
||||
STD_TORCH_CHECK(b_zeros.value().device().is_cuda(),
|
||||
"b_zeros is not on GPU");
|
||||
STD_TORCH_CHECK(b_zeros.value().is_contiguous(),
|
||||
"b_zeros is not contiguous");
|
||||
|
||||
TORCH_CHECK(b_zeros_reorder.value().device().is_cuda(),
|
||||
"b_zeros_reorder is not on GPU");
|
||||
TORCH_CHECK(b_zeros_reorder.value().is_contiguous(),
|
||||
"b_zeros_reorder is not contiguous");
|
||||
STD_TORCH_CHECK(b_zeros_reorder.value().device().is_cuda(),
|
||||
"b_zeros_reorder is not on GPU");
|
||||
STD_TORCH_CHECK(b_zeros_reorder.value().is_contiguous(),
|
||||
"b_zeros_reorder is not contiguous");
|
||||
}
|
||||
|
||||
const uint8_t* matB = reinterpret_cast<const uint8_t*>(b_qweight.data_ptr());
|
||||
@@ -136,18 +145,20 @@ void rearrange_kn_weight_as_n32k16_order(
|
||||
const void* b_zero = has_zp ? b_zeros.value().data_ptr() : nullptr;
|
||||
|
||||
uint8_t* matB_reorder =
|
||||
reinterpret_cast<uint8_t*>(b_qweight_reorder.data_ptr());
|
||||
void* b_scale_reorder = b_scales_reorder.data_ptr();
|
||||
void* b_zero_reorder = has_zp ? b_zeros_reorder.value().data_ptr() : nullptr;
|
||||
reinterpret_cast<uint8_t*>(b_qweight_reorder.mutable_data_ptr());
|
||||
void* b_scale_reorder = b_scales_reorder.mutable_data_ptr();
|
||||
void* b_zero_reorder =
|
||||
has_zp ? b_zeros_reorder.value().mutable_data_ptr() : nullptr;
|
||||
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
if (b_scales.dtype() == at::ScalarType::Half) {
|
||||
cudaStream_t stream = get_current_cuda_stream();
|
||||
if (b_scales.scalar_type() == torch::headeronly::ScalarType::Half) {
|
||||
allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__half>(
|
||||
matB, reinterpret_cast<const __half*>(b_scale),
|
||||
reinterpret_cast<const __half*>(b_zero), matB_reorder,
|
||||
reinterpret_cast<__half*>(b_scale_reorder),
|
||||
reinterpret_cast<__half*>(b_zero_reorder), K, N, N_32align, stream);
|
||||
} else if (b_scales.dtype() == at::ScalarType::BFloat16) {
|
||||
} else if (b_scales.scalar_type() ==
|
||||
torch::headeronly::ScalarType::BFloat16) {
|
||||
allspark::rearrange_kn_weight_as_n32k16_order_ldg16<__nv_bfloat16>(
|
||||
matB, reinterpret_cast<const __nv_bfloat16*>(b_scale),
|
||||
reinterpret_cast<const __nv_bfloat16*>(b_zero), matB_reorder,
|
||||
@@ -157,7 +168,7 @@ void rearrange_kn_weight_as_n32k16_order(
|
||||
}
|
||||
}
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
|
||||
m.impl("rearrange_kn_weight_as_n32k16_order",
|
||||
&rearrange_kn_weight_as_n32k16_order);
|
||||
TORCH_BOX(&rearrange_kn_weight_as_n32k16_order));
|
||||
}
|
||||
+5
-6
@@ -1,13 +1,12 @@
|
||||
#pragma once
|
||||
|
||||
#include <torch/all.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <iostream>
|
||||
#include "../marlin/marlin_dtypes.cuh"
|
||||
|
||||
#include "quantization/marlin/marlin_dtypes.cuh"
|
||||
using marlin::MarlinScalarType2;
|
||||
|
||||
namespace allspark {
|
||||
+41
-43
@@ -11,18 +11,16 @@ Redistribution and use in source and binary forms, with or without modification,
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
***********/
|
||||
|
||||
#include <torch/all.h>
|
||||
#include "libtorch_stable/torch_utils.h"
|
||||
#include "libtorch_stable/dispatch_utils.h"
|
||||
|
||||
#include <torch/csrc/stable/library.h>
|
||||
#include <torch/csrc/stable/macros.h>
|
||||
|
||||
#include <stdint.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <mma.h>
|
||||
#include <cuda/annotated_ptr>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
#include "core/registration.h"
|
||||
#include "dispatch_utils.h"
|
||||
|
||||
namespace hadacore {
|
||||
|
||||
@@ -65,12 +63,12 @@ constexpr int launch_configs_big[7][3] = {
|
||||
};
|
||||
|
||||
// a 4x2, b 2x2, c 2x2
|
||||
template <torch::ScalarType dtype>
|
||||
template <torch::headeronly::ScalarType dtype>
|
||||
__device__ __forceinline__ void mma_m16_n8_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32& c0, b32& c1){
|
||||
static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16);
|
||||
static_assert(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16);
|
||||
// d, a, b, c
|
||||
b32 zero = 0;
|
||||
if constexpr(dtype == torch::ScalarType::Half) {
|
||||
if constexpr(dtype == torch::headeronly::ScalarType::Half) {
|
||||
asm (
|
||||
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
|
||||
"{%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%8, %9};\n\t"
|
||||
@@ -89,7 +87,7 @@ __device__ __forceinline__ void mma_m16_n8_k16_b16_b16_b16_noacc(b32 a0, b32 a1,
|
||||
}
|
||||
|
||||
// a 4x2, b 4x2, c 4x2
|
||||
template <torch::ScalarType dtype>
|
||||
template <torch::headeronly::ScalarType dtype>
|
||||
__device__ __forceinline__ void mma_m16_n16_k16_b16_b16_b16_noacc(b32 a0, b32 a1, b32 a2, b32 a3, b32 b0, b32 b1, b32 b2, b32 b3, b32& c0, b32& c1, b32& c2, b32& c3){
|
||||
mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b0, b1, c0, c1);
|
||||
mma_m16_n8_k16_b16_b16_b16_noacc<dtype>(a0, a1, a2, a3, b2, b3, c2, c3);
|
||||
@@ -108,11 +106,11 @@ __device__ __forceinline__ void matrix_transpose_m8_n8_b16_inplace(b32& a0) {
|
||||
#define n_p(i) ((val_1n[i] & 0x0000FFFF) | val_1p[i] << 16)
|
||||
#define n_n(i) ((val_1n[i] & 0x0000FFFF) | val_1n[i] << 16)
|
||||
|
||||
template<int64_t num_chunks, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool enable_mask, torch::ScalarType dtype>
|
||||
template<int64_t num_chunks, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool enable_mask, torch::headeronly::ScalarType dtype>
|
||||
__global__ void __launch_bounds__(32 * warps_per_block, blocks_per_sm)
|
||||
// a is column major, b is row major
|
||||
hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) {
|
||||
static_assert(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
|
||||
static_assert(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
|
||||
|
||||
b32 b_frag_all[num_chunks][4]; // for all chunks, holds matrix fragment (which takes 4 regs of b16x2 * 32 threads)
|
||||
|
||||
@@ -162,8 +160,8 @@ hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) {
|
||||
constexpr b16 bf16_1p[4] = {0b0011111100110101, 0b0011111100000000, 0b0011111010110101, 0b0011111010000000};
|
||||
constexpr b16 bf16_1n[4] = {0b1011111100110101, 0b1011111100000000, 0b1011111010110101, 0b1011111010000000};
|
||||
|
||||
#define val_type_1p(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1p[i]) : (bf16_1p[i]))
|
||||
#define val_type_1n(i) (((dtype) == torch::ScalarType::Half) ? (fp16_1n[i]) : (bf16_1n[i]))
|
||||
#define val_type_1p(i) (((dtype) == torch::headeronly::ScalarType::Half) ? (fp16_1p[i]) : (bf16_1p[i]))
|
||||
#define val_type_1n(i) (((dtype) == torch::headeronly::ScalarType::Half) ? (fp16_1n[i]) : (bf16_1n[i]))
|
||||
constexpr b16 val_1p[4] = {val_type_1p(0), val_type_1p(1), val_type_1p(2), val_type_1p(3)};
|
||||
constexpr b16 val_1n[4] = {val_type_1n(0), val_type_1n(1), val_type_1n(2), val_type_1n(3)};
|
||||
|
||||
@@ -684,14 +682,14 @@ constexpr int64_t ceil_div(int64_t a, int64_t b) {
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
|
||||
template <torch::ScalarType dtype, int64_t chunks_per_warp, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool check_masking = false>
|
||||
template <torch::headeronly::ScalarType dtype, int64_t chunks_per_warp, int64_t warps_per_block, int64_t log_had_size, int64_t blocks_per_sm, bool check_masking = false>
|
||||
void __forceinline__ run_kernel(b16* a_mat, b16* out, int64_t num_chunks, cudaStream_t stream) {
|
||||
int64_t shared_size = chunks_per_warp * warps_per_block * 128 * 4;
|
||||
dim3 block_size = 32 * warps_per_block;
|
||||
|
||||
#define CHECK_SHARED_LIM() { \
|
||||
if (shared_size > 48 * 1024) { \
|
||||
C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536)); \
|
||||
STD_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 65536)); \
|
||||
} \
|
||||
} \
|
||||
|
||||
@@ -714,10 +712,10 @@ void __forceinline__ run_kernel(b16* a_mat, b16* out, int64_t num_chunks, cudaSt
|
||||
kernel<<<dim3(grid_size), dim3(block_size), shared_size, stream>>>(a_mat, out, num_chunks);
|
||||
}
|
||||
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
STD_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
template <torch::ScalarType dtype>
|
||||
template <torch::headeronly::ScalarType dtype>
|
||||
void run_fht(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream) {
|
||||
int64_t num_chunks = numel / 256; // caller required to ensure divisible by 256
|
||||
// for size 256, use (2, 1)
|
||||
@@ -764,54 +762,54 @@ void run_fht(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cu
|
||||
}
|
||||
}
|
||||
|
||||
template void run_fht<torch::ScalarType::Half>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
|
||||
template void run_fht<torch::ScalarType::BFloat16>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
|
||||
template void run_fht<torch::headeronly::ScalarType::Half>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
|
||||
template void run_fht<torch::headeronly::ScalarType::BFloat16>(void* a_mat_ptr, void* out_ptr, int64_t numel, int64_t had_size, cudaStream_t stream);
|
||||
|
||||
} // namespace hadacore
|
||||
|
||||
constexpr bool is_power_of_two(int x) { return x && !(x & (x - 1)); }
|
||||
|
||||
torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace) {
|
||||
torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, bool inplace) {
|
||||
auto dtype = x.scalar_type();
|
||||
TORCH_CHECK(dtype == torch::ScalarType::Half || dtype == torch::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
|
||||
TORCH_CHECK(x.is_cuda());
|
||||
|
||||
STD_TORCH_CHECK(dtype == torch::headeronly::ScalarType::Half || dtype == torch::headeronly::ScalarType::BFloat16, "Only fp16 and bf16 supported currently");
|
||||
STD_TORCH_CHECK(x.is_cuda());
|
||||
|
||||
const int had_size = x.size(-1);
|
||||
TORCH_CHECK(is_power_of_two(had_size) && (had_size <= (1U << 15)),
|
||||
STD_TORCH_CHECK(is_power_of_two(had_size) && (had_size <= (1U << 15)),
|
||||
"Only power of two Hadamard sizes up to 2^15 are supported, got ", had_size);
|
||||
|
||||
|
||||
const auto res_shape = x.sizes();
|
||||
x = x.reshape({-1, had_size});
|
||||
|
||||
x = torch::stable::reshape(x, {-1, had_size});
|
||||
|
||||
auto numel = x.numel();
|
||||
if (numel % 256 != 0) {
|
||||
x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, 0, 0, (256 - numel % 256) / had_size}));
|
||||
x = torch::stable::pad(x, {0, 0, 0, (256 - numel % 256) / had_size});
|
||||
}
|
||||
|
||||
|
||||
if (x.stride(-1) != 1) {
|
||||
x = x.contiguous();
|
||||
x = torch::stable::contiguous(x);
|
||||
}
|
||||
torch::Tensor out = inplace ? x : torch::empty_like(x);
|
||||
torch::stable::Tensor out = inplace ? x : torch::stable::empty_like(x);
|
||||
|
||||
at::cuda::CUDAGuard device_guard{(char)x.get_device()};
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
torch::stable::accelerator::DeviceGuard device_guard(x.get_device_index());
|
||||
auto stream = get_current_cuda_stream();
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(x.scalar_type(), "hadacore_transform_runfht", [&] {
|
||||
auto constexpr SCALAR_TYPE = c10::CppTypeToScalarType<scalar_t>::value;
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(x.scalar_type(), "hadacore_transform_runfht", [&] {
|
||||
auto constexpr SCALAR_TYPE = torch::headeronly::CppTypeToScalarType<scalar_t>::value;
|
||||
hadacore::run_fht<SCALAR_TYPE>(x.data_ptr(), x.data_ptr(), x.numel(), had_size, stream);
|
||||
});
|
||||
|
||||
if (numel % 256 != 0) {
|
||||
out = out.narrow(0, 0, numel / had_size);
|
||||
out = torch::stable::narrow(out, 0, 0, numel / had_size);
|
||||
}
|
||||
|
||||
if (inplace && out.data_ptr() != x.data_ptr()) {
|
||||
x.copy_(out.view(res_shape));
|
||||
torch::stable::copy_(x, torch::stable::view(out, res_shape));
|
||||
return x;
|
||||
}
|
||||
return out.reshape(res_shape);
|
||||
return torch::stable::reshape(out, res_shape);
|
||||
}
|
||||
|
||||
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
|
||||
m.impl("hadacore_transform", &hadacore_transform);
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) {
|
||||
m.impl("hadacore_transform", TORCH_BOX(&hadacore_transform));
|
||||
}
|
||||
@@ -156,6 +156,17 @@ 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,
|
||||
@@ -247,11 +258,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>
|
||||
template <typename T, typename DST_DTYPE, int GROUP_SIZE, int kGroupsPerBlockX,
|
||||
int kRowsPerBlock>
|
||||
__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 int64_t num_groups_padded,
|
||||
const int groups_per_block, const int padded_groups_per_row,
|
||||
unsigned int* __restrict__ output_s_packed, 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) {
|
||||
@@ -260,27 +271,25 @@ __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 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) {
|
||||
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) {
|
||||
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
|
||||
@@ -443,34 +452,53 @@ 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;
|
||||
const int groups_per_block = GetGroupsPerBlock(num_groups_padded);
|
||||
|
||||
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, ").");
|
||||
|
||||
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(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)); \
|
||||
#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); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
@@ -488,6 +516,7 @@ 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,
|
||||
|
||||
@@ -218,7 +218,54 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
|
||||
ops.def(
|
||||
"cutlass_encode_and_reorder_int4b_grouped(Tensor b_tensors) -> (Tensor, "
|
||||
"Tensor)");
|
||||
|
||||
// SM100 CUTLASS MLA decode
|
||||
// conditionally compiled so impl registrations are in source file
|
||||
ops.def(
|
||||
"sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
|
||||
" Tensor q_pe, Tensor kv_c_and_k_pe_cache,"
|
||||
" Tensor seq_lens, Tensor page_table,"
|
||||
" Tensor workspace, float scale,"
|
||||
" int num_kv_splits) -> ()");
|
||||
|
||||
ops.def(
|
||||
"sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches,"
|
||||
" int sm_count, int num_kv_splits) "
|
||||
"-> int");
|
||||
// Quantized GEMM for AWQ.
|
||||
ops.def(
|
||||
"awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
|
||||
"Tensor _zeros, SymInt split_k_iters) -> Tensor");
|
||||
|
||||
// Dequantization for AWQ.
|
||||
ops.def(
|
||||
"awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
|
||||
"Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
|
||||
|
||||
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
|
||||
// conditionally compiled so impl registration is in source file
|
||||
ops.def(
|
||||
"dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
|
||||
|
||||
// reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
|
||||
ops.def(
|
||||
"rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
|
||||
"Tensor? b_zeros, "
|
||||
"bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, "
|
||||
"Tensor!? b_zeros_reorder, "
|
||||
"int K, int N, int N_32align) -> ()");
|
||||
|
||||
// AllSpark quantization ops
|
||||
ops.def(
|
||||
"allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, "
|
||||
"Tensor? b_qzeros, "
|
||||
"SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
|
||||
"CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
|
||||
#endif
|
||||
|
||||
// Hadamard transforms
|
||||
// conditionally compiled so impl registration is in source file
|
||||
ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
|
||||
}
|
||||
|
||||
STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
@@ -254,6 +301,16 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
|
||||
ops.impl("silu_and_mul_nvfp4_quant", TORCH_BOX(&silu_and_mul_nvfp4_quant));
|
||||
// mxfp4_experts_quant: registered in mxfp4_experts_quant.cu (SM100 only).
|
||||
// W4A8 ops: registered in w4a8_mm_entry.cu / w4a8_grouped_mm_entry.cu.
|
||||
|
||||
// AWQ ops
|
||||
ops.impl("awq_gemm", TORCH_BOX(&awq_gemm));
|
||||
ops.impl("awq_dequantize", TORCH_BOX(&awq_dequantize));
|
||||
|
||||
// DSV3 fused A GEMM: conditionally compiled so impl registration is in
|
||||
// source file (dsv3_fused_a_gemm.cu)
|
||||
|
||||
// AllSpark ops: conditionally compiled so impl registrations are in source
|
||||
// files (allspark_repack.cu and allspark_qgemm_w8a16.cu)
|
||||
#endif
|
||||
}
|
||||
|
||||
|
||||
@@ -6,12 +6,71 @@
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/headeronly/util/shim_utils.h>
|
||||
|
||||
#include <cublas_v2.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <deque>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
// Stable ABI equivalent of TORCH_CHECK_NOT_IMPLEMENTED.
|
||||
#define STD_TORCH_CHECK_NOT_IMPLEMENTED(cond, ...) \
|
||||
STD_TORCH_CHECK(cond, "NotImplementedError: ", __VA_ARGS__)
|
||||
|
||||
// Device properties cache for stable ABI compatibility.
|
||||
// Uses raw CUDA/HIP APIs instead of ATen functions.
|
||||
// Using inline ensures a single instance across all translation units.
|
||||
inline std::deque<std::once_flag> device_flags;
|
||||
inline std::vector<cudaDeviceProp> device_properties;
|
||||
inline std::once_flag vectors_init_flag;
|
||||
|
||||
inline void do_init_device_vectors() {
|
||||
int device_count;
|
||||
cudaError_t err = cudaGetDeviceCount(&device_count);
|
||||
if (err != cudaSuccess) {
|
||||
STD_TORCH_CHECK(false, "cudaGetDeviceCount failed: " +
|
||||
std::string(cudaGetErrorString(err)));
|
||||
}
|
||||
device_flags.resize(device_count);
|
||||
device_properties.resize(device_count);
|
||||
}
|
||||
|
||||
inline void initDeviceVectors() {
|
||||
std::call_once(vectors_init_flag, do_init_device_vectors);
|
||||
}
|
||||
|
||||
inline void initDeviceProperty(int device_index) {
|
||||
cudaDeviceProp device_prop{};
|
||||
cudaError_t err = cudaGetDeviceProperties(&device_prop, device_index);
|
||||
if (err != cudaSuccess) {
|
||||
STD_TORCH_CHECK(false, "cudaGetDeviceProperties failed: " +
|
||||
std::string(cudaGetErrorString(err)));
|
||||
}
|
||||
device_properties[device_index] = device_prop;
|
||||
}
|
||||
|
||||
// Get device properties using raw CUDA/HIP APIs (stable ABI compatible).
|
||||
// Caches results per device so cudaGetDeviceProperties is called at most once
|
||||
// per device.
|
||||
inline cudaDeviceProp* get_device_prop() {
|
||||
initDeviceVectors();
|
||||
int device_index;
|
||||
cudaError_t err = cudaGetDevice(&device_index);
|
||||
if (err != cudaSuccess) {
|
||||
STD_TORCH_CHECK(
|
||||
false, "cudaGetDevice failed: " + std::string(cudaGetErrorString(err)));
|
||||
}
|
||||
STD_TORCH_CHECK(device_index >= 0 && static_cast<size_t>(device_index) <
|
||||
device_properties.size(),
|
||||
"CUDA device index " + std::to_string(device_index) +
|
||||
" out of range [0, " +
|
||||
std::to_string(device_properties.size()) + ")");
|
||||
|
||||
std::call_once(device_flags[device_index], initDeviceProperty, device_index);
|
||||
return &device_properties[device_index];
|
||||
}
|
||||
|
||||
// Utility to get the current CUDA stream for a given device using stable APIs.
|
||||
// Returns a cudaStream_t for use in kernel launches.
|
||||
inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) {
|
||||
@@ -20,3 +79,10 @@ inline cudaStream_t get_current_cuda_stream(int32_t device_index = -1) {
|
||||
aoti_torch_get_current_cuda_stream(device_index, &stream_ptr));
|
||||
return reinterpret_cast<cudaStream_t>(stream_ptr);
|
||||
}
|
||||
|
||||
// Utility to get the current cuBLAS handle using stable APIs.
|
||||
inline cublasHandle_t get_current_cuda_blas_handle() {
|
||||
void* blas_handle_ptr = nullptr;
|
||||
TORCH_ERROR_CODE_CHECK(torch_get_current_cuda_blas_handle(&blas_handle_ptr));
|
||||
return reinterpret_cast<cublasHandle_t>(blas_handle_ptr);
|
||||
}
|
||||
|
||||
-20
@@ -200,19 +200,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
|
||||
|
||||
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
|
||||
torch::Tensor awq_gemm(torch::Tensor _in_feats, torch::Tensor _kernel,
|
||||
torch::Tensor _scaling_factors, torch::Tensor _zeros,
|
||||
int64_t split_k_iters);
|
||||
|
||||
torch::Tensor awq_dequantize(torch::Tensor _kernel,
|
||||
torch::Tensor _scaling_factors,
|
||||
torch::Tensor _zeros, int64_t split_k_iters,
|
||||
int64_t thx, int64_t thy);
|
||||
|
||||
#endif
|
||||
|
||||
torch::Tensor ggml_dequantize(torch::Tensor W, int64_t type, int64_t m,
|
||||
int64_t n,
|
||||
std::optional<at::ScalarType> const& dtype);
|
||||
@@ -302,8 +289,6 @@ std::tuple<int64_t, torch::Tensor> allocate_shared_buffer_and_handle(
|
||||
int64_t open_mem_handle(torch::Tensor& mem_handle);
|
||||
void free_shared_buffer(int64_t buffer);
|
||||
|
||||
torch::Tensor hadacore_transform(torch::Tensor& x, bool inplace);
|
||||
|
||||
#ifdef USE_ROCM
|
||||
fptr_t init_custom_qr(int64_t rank, int64_t world_size,
|
||||
std::optional<int64_t> qr_max_size = std::nullopt);
|
||||
@@ -315,11 +300,6 @@ void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
|
||||
int64_t qr_max_size();
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
|
||||
torch::Tensor const& mat_b);
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
torch::Tensor minimax_allreduce_rms(torch::Tensor const& input,
|
||||
torch::Tensor const& norm_weight,
|
||||
|
||||
@@ -2,10 +2,14 @@
|
||||
|
||||
#ifndef _marlin_cuh
|
||||
#define _marlin_cuh
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
// These torch headers are only needed by non-stable callers (e.g. ops.cu).
|
||||
// Guard them so that stable ABI targets can still include marlin.cuh
|
||||
// for Vec, constants, and cp_async helpers without pulling in torch/all.h.
|
||||
#ifndef TORCH_TARGET_VERSION
|
||||
#include <torch/all.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#endif
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
@@ -263,22 +263,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
|
||||
// Quantization ops
|
||||
#ifndef USE_ROCM
|
||||
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
|
||||
ops.def(
|
||||
"dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
|
||||
// conditionally compiled so impl registration is in source file
|
||||
|
||||
// Quantized GEMM for AWQ.
|
||||
ops.def(
|
||||
"awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
|
||||
"Tensor _zeros, SymInt split_k_iters) -> Tensor");
|
||||
ops.impl("awq_gemm", torch::kCUDA, &awq_gemm);
|
||||
|
||||
// Dequantization for AWQ.
|
||||
ops.def(
|
||||
"awq_dequantize(Tensor _kernel, Tensor _scaling_factors, "
|
||||
"Tensor _zeros, SymInt split_k_iters, int thx, int thy) -> Tensor");
|
||||
ops.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);
|
||||
|
||||
// Note about marlin kernel 'workspace' arguments:
|
||||
// Technically these should be mutable since they are modified by the kernel.
|
||||
@@ -408,22 +392,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
" -> ()");
|
||||
// conditionally compiled so impl registration is in source file
|
||||
|
||||
// SM100 CUTLASS MLA decode
|
||||
ops.def(
|
||||
"sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
|
||||
" Tensor q_pe, Tensor kv_c_and_k_pe_cache,"
|
||||
" Tensor seq_lens, Tensor page_table,"
|
||||
" Tensor workspace, float scale,"
|
||||
" int num_kv_splits) -> ()");
|
||||
// conditionally compiled so impl in source file
|
||||
|
||||
// SM100 CUTLASS MLA workspace
|
||||
ops.def(
|
||||
"sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches,"
|
||||
" int sm_count, int num_kv_splits) "
|
||||
"-> int");
|
||||
// conditionally compiled so impl in source file
|
||||
|
||||
#endif
|
||||
|
||||
// Quantized GEMM for GPTQ.
|
||||
@@ -496,26 +464,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"Tensor? last_chunk_indices) -> ()");
|
||||
ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
|
||||
|
||||
// Hadamard transforms
|
||||
ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor");
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
|
||||
ops.def(
|
||||
"rearrange_kn_weight_as_n32k16_order(Tensor b_qweight, Tensor b_scales, "
|
||||
"Tensor? b_zeros, "
|
||||
"bool has_zp, Tensor! b_qweight_reorder, Tensor! b_scales_reorder, "
|
||||
"Tensor!? b_zeros_reorder, "
|
||||
"int K, int N, int N_32align) -> ()");
|
||||
// conditionally compiled so impl in source file
|
||||
|
||||
// AllSpark quantization ops
|
||||
ops.def(
|
||||
"allspark_w8a16_gemm(Tensor a, Tensor b_qweight, Tensor b_scales, "
|
||||
"Tensor? b_qzeros, "
|
||||
"SymInt n, SymInt group_size, SymInt sm_count, SymInt sm_version, SymInt "
|
||||
"CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor");
|
||||
|
||||
ops.def(
|
||||
"minimax_allreduce_rms("
|
||||
"Tensor input,"
|
||||
|
||||
+17
-1
@@ -199,7 +199,10 @@ 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 [ "${PYTORCH_NIGHTLY}" = "1" ]; then \
|
||||
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 \
|
||||
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 '.') \
|
||||
@@ -301,6 +304,14 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
python3 use_existing_torch.py --prefix; \
|
||||
fi
|
||||
|
||||
# Provision one bare Python per `requires-python` entry; cmake reads
|
||||
# DEEPGEMM_PYTHON_INTERPRETERS to build DeepGEMM `_C` for each. See
|
||||
# cmake/external_projects/deepgemm.cmake for the full picture.
|
||||
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)
|
||||
@@ -328,6 +339,7 @@ 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; \
|
||||
@@ -345,6 +357,7 @@ 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
|
||||
|
||||
@@ -616,6 +629,9 @@ 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
|
||||
|
||||
@@ -34,6 +34,7 @@ you may contact the following individuals:
|
||||
|
||||
- Simon Mo - <simon.mo@hey.com>
|
||||
- Russell Bryant - <rbryant@redhat.com>
|
||||
- Juan Pérez de Algaba - <jperezde@redhat.com>
|
||||
- Huzaifa Sidhpurwala - <huzaifas@redhat.com>
|
||||
|
||||
## Slack Discussion
|
||||
|
||||
@@ -2,4 +2,4 @@
|
||||
|
||||
vLLM can be deployed with [llm-d](https://github.com/llm-d/llm-d), a Kubernetes-native distributed inference serving stack providing well-lit paths for anyone to serve large generative AI models at scale. It helps achieve the fastest "time to state-of-the-art (SOTA) performance" for key OSS models across most hardware accelerators and infrastructure providers.
|
||||
|
||||
You can use vLLM with llm-d directly by following [this guide](https://llm-d.ai/docs/guide) or via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview).
|
||||
You can use vLLM with llm-d directly by following [the official guides](https://llm-d.ai/docs/guides) or via [KServe's LLMInferenceService](https://kserve.github.io/website/docs/model-serving/generative-inference/llmisvc/llmisvc-overview).
|
||||
|
||||
@@ -85,8 +85,9 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
|
||||
|
||||
| Architecture | Models | CG for Image | CG for Video |
|
||||
| ------------ | ------ | ------------ | ------------ |
|
||||
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
|
||||
| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
|
||||
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
|
||||
| `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ |
|
||||
|
||||
!!! note
|
||||
Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
|
||||
|
||||
@@ -81,7 +81,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
|
||||
| Kernel | Input act. format | Quant. types | Quant. format | Activation function | Apply Weight On Input | Modular | Source |
|
||||
| ------ | ----------------- | ------------ | ------------- | ------------------- | --------------------- | ------- | ------ |
|
||||
| triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.experts.triton_moe.TritonExperts] |
|
||||
| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
|
||||
| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.BatchedTritonExperts] |
|
||||
| deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | </br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
|
||||
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp4] |
|
||||
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassBatchedExpertsFp8] |
|
||||
@@ -91,7 +91,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
|
||||
| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
|
||||
| rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |
|
||||
| cpu_fused_moe | standard | N/A | N/A | silu | N | N | [`CPUFusedMOE`][vllm.model_executor.layers.fused_moe.cpu_fused_moe.CPUFusedMOE] |
|
||||
| naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.NaiveBatchedExperts] |
|
||||
| naive batched<sup>4</sup> | batched | int8,</br>fp8 | G,A,T | silu, gelu | <sup>6</sup> | Y | [`NaiveBatchedExperts`][vllm.model_executor.layers.fused_moe.experts.fused_batched_moe.NaiveBatchedExperts] |
|
||||
|
||||
!!! info "Table key"
|
||||
1. All types: mxfp4, nvfp4, int4, int8, fp8
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
# 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.
|
||||
@@ -2179,6 +2179,92 @@ def run_qwen3_vl_moe(questions: list[str], modality: str) -> ModelRequestData:
|
||||
)
|
||||
|
||||
|
||||
# Qwen3.5-Dense
|
||||
def run_qwen3_5(questions: list[str], modality: str) -> ModelRequestData:
|
||||
model_name = "Qwen/Qwen3.5-4B"
|
||||
|
||||
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_model_len=4096,
|
||||
max_num_seqs=5,
|
||||
mm_processor_kwargs={
|
||||
"min_pixels": 28 * 28,
|
||||
"max_pixels": 1280 * 28 * 28,
|
||||
"fps": 1,
|
||||
},
|
||||
limit_mm_per_prompt=mm_limit,
|
||||
)
|
||||
|
||||
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
|
||||
|
||||
if modality == "image":
|
||||
placeholder = image_placeholder
|
||||
elif modality == "video":
|
||||
placeholder = video_placeholder
|
||||
elif modality == "image+video":
|
||||
placeholder = image_placeholder + video_placeholder
|
||||
|
||||
prompts = [
|
||||
(
|
||||
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
|
||||
f"<|im_start|>user\n{placeholder}"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
for question in questions
|
||||
]
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompts=prompts,
|
||||
)
|
||||
|
||||
|
||||
# Qwen3.5-MoE
|
||||
def run_qwen3_5_moe(questions: list[str], modality: str) -> ModelRequestData:
|
||||
model_name = "Qwen/Qwen3.5-35B-A3B"
|
||||
|
||||
mm_limit = {"image": 1, "video": 1} if modality == "image+video" else {modality: 1}
|
||||
engine_args = EngineArgs(
|
||||
model=model_name,
|
||||
max_model_len=4096,
|
||||
max_num_seqs=5,
|
||||
mm_processor_kwargs={
|
||||
"min_pixels": 28 * 28,
|
||||
"max_pixels": 1280 * 28 * 28,
|
||||
"fps": 1,
|
||||
},
|
||||
limit_mm_per_prompt=mm_limit,
|
||||
)
|
||||
|
||||
image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
video_placeholder = "<|vision_start|><|video_pad|><|vision_end|>"
|
||||
|
||||
if modality == "image":
|
||||
placeholder = image_placeholder
|
||||
elif modality == "video":
|
||||
placeholder = video_placeholder
|
||||
elif modality == "image+video":
|
||||
placeholder = image_placeholder + video_placeholder
|
||||
|
||||
prompts = [
|
||||
(
|
||||
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
|
||||
f"<|im_start|>user\n{placeholder}"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
for question in questions
|
||||
]
|
||||
|
||||
return ModelRequestData(
|
||||
engine_args=engine_args,
|
||||
prompts=prompts,
|
||||
)
|
||||
|
||||
|
||||
# R-4B
|
||||
def run_r_vl(questions: list[str], modality: str) -> ModelRequestData:
|
||||
assert modality == "image"
|
||||
@@ -2442,6 +2528,8 @@ model_example_map = {
|
||||
"qwen2_5_omni": run_qwen2_5_omni,
|
||||
"qwen3_vl": run_qwen3_vl,
|
||||
"qwen3_vl_moe": run_qwen3_vl_moe,
|
||||
"qwen3_5": run_qwen3_5,
|
||||
"qwen3_5_moe": run_qwen3_5_moe,
|
||||
"rvl": run_r_vl,
|
||||
"skywork_chat": run_skyworkr1v,
|
||||
"smolvlm": run_smolvlm,
|
||||
@@ -2460,13 +2548,17 @@ MODELS_NEED_VIDEO_METADATA = [
|
||||
"molmo2",
|
||||
"qwen3_vl",
|
||||
"qwen3_vl_moe",
|
||||
"qwen3_5",
|
||||
"qwen3_5_moe",
|
||||
]
|
||||
|
||||
|
||||
MODELS_SUPPORT_VIT_CUDA_GRAPH = [
|
||||
"qwen2_5_vl",
|
||||
"qwen3_vl",
|
||||
"qwen3_vl_moe",
|
||||
"qwen2_5_vl",
|
||||
"qwen3_5",
|
||||
"qwen3_5_moe",
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -20,6 +20,6 @@ nvidia-cudnn-frontend>=1.13.0,<1.19.0
|
||||
# Required for faster safetensors model loading
|
||||
fastsafetensors >= 0.2.2
|
||||
|
||||
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
|
||||
nvidia-cutlass-dsl>=4.4.2
|
||||
# QuACK and Cutlass DSL for FA4 and DeepSeek-V4 kernels
|
||||
nvidia-cutlass-dsl[cu13]>=4.5.0
|
||||
quack-kernels>=0.3.3
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
lmcache >= 0.3.9
|
||||
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
|
||||
nixl >= 1.1.0 # Required for disaggregated prefill
|
||||
mooncake-transfer-engine >= 0.3.8
|
||||
|
||||
@@ -15,4 +15,4 @@ torch==2.11.0+xpu
|
||||
torchaudio
|
||||
torchvision
|
||||
|
||||
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.7/vllm_xpu_kernels-0.1.7-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.8/vllm_xpu_kernels-0.1.8-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
|
||||
@@ -970,6 +970,9 @@ 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():
|
||||
|
||||
@@ -1686,3 +1686,44 @@ def disable_log_dedup(monkeypatch):
|
||||
logger._print_warning_once = original_print_warning_once
|
||||
logger._print_info_once = original_print_info_once
|
||||
logger._print_debug_once = original_print_debug_once
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def fake_vllm_ir(monkeypatch):
|
||||
"""
|
||||
Pytest fixture to allow isolated IR op registration in tests.
|
||||
|
||||
Replaces IrOp.registry with an empty dict and swaps ``vllm_ir_torch_lib`` for a
|
||||
fresh ``Library`` with a unique namespace per test (see ``Library.ns``).
|
||||
|
||||
Torch keeps registrations for the process lifetime; reusing the fragment
|
||||
name ``vllm_ir`` and defining the same op string again can segfault. A
|
||||
random library name keeps each fixture run on a disjoint namespace.
|
||||
|
||||
The test Library is kept alive until after monkeypatch teardown so PyTorch's
|
||||
C++ state is not freed while references may still exist.
|
||||
|
||||
Usage:
|
||||
def test_my_ir_op(fake_vllm_ir):
|
||||
@vllm.ir.register_op
|
||||
def my_test_op(x: torch.Tensor) -> torch.Tensor:
|
||||
return x * 2
|
||||
|
||||
result = my_test_op(torch.tensor([1, 2, 3]))
|
||||
# Registry and library cleaned up automatically after the test
|
||||
"""
|
||||
import secrets
|
||||
|
||||
from torch.library import Library
|
||||
from vllm.ir.op import IrOp
|
||||
|
||||
monkeypatch.setattr(IrOp, "registry", {})
|
||||
|
||||
# Keep a local reference so the Library is not GC'd before monkeypatch
|
||||
# teardown restores the original reference.
|
||||
test_lib = Library(f"vllm_ir_{secrets.token_hex(8)}", "FRAGMENT")
|
||||
monkeypatch.setattr("vllm.ir.op.vllm_ir_torch_lib", test_lib)
|
||||
|
||||
yield
|
||||
|
||||
del test_lib
|
||||
|
||||
@@ -10,6 +10,7 @@ import torch
|
||||
|
||||
from tests.kernels.moe.utils import make_test_quant_config
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed.eplb.eplb_state import EplbLayerState
|
||||
from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
|
||||
from vllm.distributed.parallel_state import (
|
||||
ensure_model_parallel_initialized,
|
||||
@@ -201,7 +202,7 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig):
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
fml.enable_eplb = True
|
||||
fml.eplb_state = EplbLayerState()
|
||||
fml.set_eplb_state(
|
||||
lidx,
|
||||
torch.zeros(
|
||||
|
||||
@@ -3,7 +3,9 @@
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.model_executor.layers.fused_moe.layer import determine_expert_map
|
||||
from vllm.model_executor.layers.fused_moe.expert_map_manager import (
|
||||
determine_expert_map,
|
||||
)
|
||||
|
||||
|
||||
def verify_round_robin_pattern(expert_map, ep_rank, ep_size, global_num_experts):
|
||||
|
||||
@@ -9,7 +9,10 @@ from vllm.distributed.kv_events import BlockRemoved, BlockStored
|
||||
_FAKE_HASH: bytes = b"\xab" * 32
|
||||
|
||||
|
||||
def _make_block_stored(group_idx: int | None = None) -> BlockStored:
|
||||
def _make_block_stored(
|
||||
group_idx: int | None = None,
|
||||
kv_cache_spec_sliding_window: int | None = None,
|
||||
) -> BlockStored:
|
||||
return BlockStored(
|
||||
block_hashes=[_FAKE_HASH],
|
||||
parent_block_hash=None,
|
||||
@@ -19,10 +22,13 @@ def _make_block_stored(group_idx: int | None = None) -> BlockStored:
|
||||
medium="GPU",
|
||||
lora_name=None,
|
||||
group_idx=group_idx,
|
||||
kv_cache_spec_sliding_window=kv_cache_spec_sliding_window,
|
||||
)
|
||||
|
||||
|
||||
def _make_block_removed(group_idx: int | None = None) -> BlockRemoved:
|
||||
def _make_block_removed(
|
||||
group_idx: int | None = None,
|
||||
) -> BlockRemoved:
|
||||
return BlockRemoved(
|
||||
block_hashes=[_FAKE_HASH],
|
||||
medium="GPU",
|
||||
@@ -72,3 +78,9 @@ def test_block_removed_hash_same_for_equal_group_idx():
|
||||
event_a = _make_block_removed(group_idx=1)
|
||||
event_b = _make_block_removed(group_idx=1)
|
||||
assert hash(event_a) == hash(event_b)
|
||||
|
||||
|
||||
def test_block_stored_hash_differs_by_sliding_window():
|
||||
event_a = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=128)
|
||||
event_b = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=256)
|
||||
assert hash(event_a) != hash(event_b)
|
||||
|
||||
@@ -4,7 +4,6 @@
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
|
||||
from openai.types.responses.response_function_tool_call_output_item import (
|
||||
ResponseFunctionToolCallOutputItem,
|
||||
@@ -17,10 +16,8 @@ from openai.types.responses.response_reasoning_item import (
|
||||
Summary,
|
||||
)
|
||||
|
||||
from vllm.entrypoints.constants import MCP_PREFIX
|
||||
from vllm.entrypoints.openai.responses.utils import (
|
||||
_construct_single_message_from_response_item,
|
||||
_maybe_combine_reasoning_and_tool_call,
|
||||
_construct_message_from_response_item,
|
||||
construct_chat_messages_with_tool_call,
|
||||
construct_input_messages,
|
||||
convert_tool_responses_to_completions_format,
|
||||
@@ -28,6 +25,96 @@ from vllm.entrypoints.openai.responses.utils import (
|
||||
)
|
||||
|
||||
|
||||
def _single_chat_message(item):
|
||||
message = _construct_message_from_response_item(item)
|
||||
assert message is not None
|
||||
return message
|
||||
|
||||
|
||||
def make_output_message(
|
||||
text: str,
|
||||
*,
|
||||
id: str = "msg_1",
|
||||
status: str = "completed",
|
||||
) -> ResponseOutputMessage:
|
||||
return ResponseOutputMessage(
|
||||
id=id,
|
||||
content=[
|
||||
ResponseOutputText(
|
||||
annotations=[],
|
||||
text=text,
|
||||
type="output_text",
|
||||
logprobs=None,
|
||||
)
|
||||
],
|
||||
role="assistant",
|
||||
status=status,
|
||||
type="message",
|
||||
)
|
||||
|
||||
|
||||
def make_reasoning_item(
|
||||
*,
|
||||
content_text: str | None = None,
|
||||
summary_text: str | None = None,
|
||||
content: list[Content] | None = None,
|
||||
summary: list[Summary] | None = None,
|
||||
encrypted_content: str | None = None,
|
||||
id: str = "reasoning_1",
|
||||
status: str | None = None,
|
||||
) -> ResponseReasoningItem:
|
||||
if content is None and content_text is not None:
|
||||
content = [Content(text=content_text, type="reasoning_text")]
|
||||
if summary is None and summary_text is not None:
|
||||
summary = [Summary(text=summary_text, type="summary_text")]
|
||||
|
||||
return ResponseReasoningItem(
|
||||
id=id,
|
||||
summary=[] if summary is None else summary,
|
||||
type="reasoning",
|
||||
content=content,
|
||||
encrypted_content=encrypted_content,
|
||||
status=status,
|
||||
)
|
||||
|
||||
|
||||
def make_function_call(
|
||||
*,
|
||||
call_id: str,
|
||||
name: str = "test_function",
|
||||
arguments: str = "{}",
|
||||
id: str = "tool_id",
|
||||
status: str | None = None,
|
||||
) -> ResponseFunctionToolCall:
|
||||
kwargs = {
|
||||
"type": "function_call",
|
||||
"id": id,
|
||||
"call_id": call_id,
|
||||
"name": name,
|
||||
"arguments": arguments,
|
||||
}
|
||||
if status is not None:
|
||||
kwargs["status"] = status
|
||||
|
||||
return ResponseFunctionToolCall(**kwargs)
|
||||
|
||||
|
||||
def make_function_call_output(
|
||||
*,
|
||||
call_id: str,
|
||||
output: str = "42",
|
||||
id: str = "output_1",
|
||||
status: str = "completed",
|
||||
) -> ResponseFunctionToolCallOutputItem:
|
||||
return ResponseFunctionToolCallOutputItem(
|
||||
id=id,
|
||||
type="function_call_output",
|
||||
call_id=call_id,
|
||||
output=output,
|
||||
status=status,
|
||||
)
|
||||
|
||||
|
||||
class TestResponsesUtils:
|
||||
"""Tests for convert_tool_responses_to_completions_format function."""
|
||||
|
||||
@@ -87,7 +174,7 @@ class TestResponsesUtils:
|
||||
message["tool_calls"][0]["function"]["arguments"] == '{"code": "123+456"}'
|
||||
)
|
||||
|
||||
def test_construct_single_message_from_response_item(self):
|
||||
def test_construct_chat_messages_preserves_single_item_conversions(self):
|
||||
item = ResponseReasoningItem(
|
||||
id="lol",
|
||||
summary=[],
|
||||
@@ -101,7 +188,7 @@ class TestResponsesUtils:
|
||||
encrypted_content=None,
|
||||
status=None,
|
||||
)
|
||||
formatted_item = _construct_single_message_from_response_item(item)
|
||||
formatted_item = _single_chat_message(item)
|
||||
assert formatted_item["role"] == "assistant"
|
||||
assert formatted_item["reasoning"] == "Leroy Jenkins"
|
||||
|
||||
@@ -119,7 +206,7 @@ class TestResponsesUtils:
|
||||
status=None,
|
||||
)
|
||||
|
||||
formatted_item = _construct_single_message_from_response_item(item)
|
||||
formatted_item = _single_chat_message(item)
|
||||
assert formatted_item["role"] == "assistant"
|
||||
assert (
|
||||
formatted_item["reasoning"]
|
||||
@@ -133,11 +220,22 @@ class TestResponsesUtils:
|
||||
output="1234",
|
||||
status="completed",
|
||||
)
|
||||
formatted_item = _construct_single_message_from_response_item(tool_call_output)
|
||||
formatted_item = _single_chat_message(tool_call_output)
|
||||
assert formatted_item["role"] == "tool"
|
||||
assert formatted_item["content"] == "1234"
|
||||
assert formatted_item["tool_call_id"] == "temp"
|
||||
|
||||
formatted_item = _single_chat_message(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "temp_dict",
|
||||
"output": "5678",
|
||||
}
|
||||
)
|
||||
assert formatted_item["role"] == "tool"
|
||||
assert formatted_item["content"] == "5678"
|
||||
assert formatted_item["tool_call_id"] == "temp_dict"
|
||||
|
||||
item = ResponseReasoningItem(
|
||||
id="lol",
|
||||
summary=[],
|
||||
@@ -147,7 +245,7 @@ class TestResponsesUtils:
|
||||
status=None,
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
_construct_single_message_from_response_item(item)
|
||||
construct_chat_messages_with_tool_call([item])
|
||||
|
||||
output_item = ResponseOutputMessage(
|
||||
id="msg_bf585bbbe3d500e0",
|
||||
@@ -164,7 +262,7 @@ class TestResponsesUtils:
|
||||
type="message",
|
||||
)
|
||||
|
||||
formatted_item = _construct_single_message_from_response_item(output_item)
|
||||
formatted_item = _single_chat_message(output_item)
|
||||
assert formatted_item["role"] == "assistant"
|
||||
assert formatted_item["content"] == "dongyi"
|
||||
|
||||
@@ -192,7 +290,7 @@ class TestReasoningItemContentPriority:
|
||||
encrypted_content=None,
|
||||
status=None,
|
||||
)
|
||||
formatted = _construct_single_message_from_response_item(item)
|
||||
formatted = _single_chat_message(item)
|
||||
assert formatted["reasoning"] == "This is the actual content"
|
||||
|
||||
def test_content_only(self):
|
||||
@@ -210,7 +308,7 @@ class TestReasoningItemContentPriority:
|
||||
encrypted_content=None,
|
||||
status=None,
|
||||
)
|
||||
formatted = _construct_single_message_from_response_item(item)
|
||||
formatted = _single_chat_message(item)
|
||||
assert formatted["reasoning"] == "Content without summary"
|
||||
|
||||
@patch("vllm.entrypoints.openai.responses.utils.logger")
|
||||
@@ -229,7 +327,7 @@ class TestReasoningItemContentPriority:
|
||||
encrypted_content=None,
|
||||
status=None,
|
||||
)
|
||||
formatted = _construct_single_message_from_response_item(item)
|
||||
formatted = _single_chat_message(item)
|
||||
assert formatted["reasoning"] == "Fallback summary text"
|
||||
mock_logger.warning.assert_called_once()
|
||||
assert (
|
||||
@@ -252,7 +350,7 @@ class TestReasoningItemContentPriority:
|
||||
encrypted_content=None,
|
||||
status=None,
|
||||
)
|
||||
formatted = _construct_single_message_from_response_item(item)
|
||||
formatted = _single_chat_message(item)
|
||||
assert formatted["reasoning"] == "Summary when content empty"
|
||||
mock_logger.warning.assert_called_once()
|
||||
assert (
|
||||
@@ -269,7 +367,7 @@ class TestReasoningItemContentPriority:
|
||||
encrypted_content=None,
|
||||
status=None,
|
||||
)
|
||||
formatted = _construct_single_message_from_response_item(item)
|
||||
formatted = _single_chat_message(item)
|
||||
assert formatted["reasoning"] == ""
|
||||
|
||||
def test_encrypted_content_raises(self):
|
||||
@@ -293,7 +391,7 @@ class TestReasoningItemContentPriority:
|
||||
status=None,
|
||||
)
|
||||
with pytest.raises(ValueError):
|
||||
_construct_single_message_from_response_item(item)
|
||||
construct_chat_messages_with_tool_call([item])
|
||||
|
||||
@patch("vllm.entrypoints.openai.responses.utils.logger")
|
||||
def test_summary_with_multiple_entries_uses_first(self, mock_logger):
|
||||
@@ -315,7 +413,7 @@ class TestReasoningItemContentPriority:
|
||||
encrypted_content=None,
|
||||
status=None,
|
||||
)
|
||||
formatted = _construct_single_message_from_response_item(item)
|
||||
formatted = _single_chat_message(item)
|
||||
assert formatted["reasoning"] == "First summary"
|
||||
mock_logger.warning.assert_called_once()
|
||||
assert (
|
||||
@@ -343,7 +441,7 @@ class TestReasoningItemContentPriority:
|
||||
encrypted_content=None,
|
||||
status=None,
|
||||
)
|
||||
_construct_single_message_from_response_item(item)
|
||||
construct_chat_messages_with_tool_call([item])
|
||||
mock_logger.warning.assert_not_called()
|
||||
|
||||
|
||||
@@ -626,119 +724,125 @@ class TestShouldContinueFinalMessage:
|
||||
assert should_continue_final_message([dict_item]) is False
|
||||
|
||||
|
||||
class TestMaybeCombineReasoningAndToolCall:
|
||||
"""Tests for _maybe_combine_reasoning_and_tool_call function."""
|
||||
class TestConstructChatMessagesCombinePolicy:
|
||||
"""Tests for contiguous assistant-side merging."""
|
||||
|
||||
def test_returns_none_when_item_id_is_none(self):
|
||||
"""
|
||||
Test fix from PR #31999: when item.id is None, should return None
|
||||
instead of raising TypeError on startswith().
|
||||
"""
|
||||
item = ResponseFunctionToolCall(
|
||||
type="function_call",
|
||||
id=None, # This was causing TypeError before the fix
|
||||
call_id="call_123",
|
||||
name="test_function",
|
||||
arguments="{}",
|
||||
)
|
||||
messages: list[ChatCompletionMessageParam] = []
|
||||
@pytest.mark.parametrize(
|
||||
("items", "expected_content", "expected_reasoning", "expected_tool_call_ids"),
|
||||
[
|
||||
pytest.param(
|
||||
[
|
||||
make_reasoning_item(content_text="Let me think"),
|
||||
make_output_message("Hello"),
|
||||
],
|
||||
"Hello",
|
||||
"Let me think",
|
||||
None,
|
||||
id="reasoning-output-messages",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
make_function_call(call_id="call_123"),
|
||||
make_function_call(call_id="call_456"),
|
||||
],
|
||||
None,
|
||||
None,
|
||||
["call_123", "call_456"],
|
||||
id="consecutive-tool-calls",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
make_reasoning_item(content_text="Let me think"),
|
||||
make_function_call(call_id="call_123"),
|
||||
],
|
||||
None,
|
||||
"Let me think",
|
||||
["call_123"],
|
||||
id="reasoning-tool-call",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
make_output_message("Hello"),
|
||||
make_function_call(call_id="call_123"),
|
||||
],
|
||||
"Hello",
|
||||
None,
|
||||
["call_123"],
|
||||
id="output-tool-call",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
make_reasoning_item(content_text="Thinking"),
|
||||
make_output_message("Hello"),
|
||||
make_function_call(call_id="call_123"),
|
||||
make_function_call(call_id="call_456"),
|
||||
],
|
||||
"Hello",
|
||||
"Thinking",
|
||||
["call_123", "call_456"],
|
||||
id="reasoning-output-tool-call",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_assistant_side_items_merge_until_tool_output(
|
||||
self,
|
||||
items,
|
||||
expected_content,
|
||||
expected_reasoning,
|
||||
expected_tool_call_ids,
|
||||
):
|
||||
messages = construct_chat_messages_with_tool_call(items)
|
||||
|
||||
result = _maybe_combine_reasoning_and_tool_call(item, messages)
|
||||
assert len(messages) == 1
|
||||
assert messages[0]["role"] == "assistant"
|
||||
if expected_content is None:
|
||||
assert "content" not in messages[0]
|
||||
else:
|
||||
assert messages[0]["content"] == expected_content
|
||||
if expected_reasoning is None:
|
||||
assert "reasoning" not in messages[0]
|
||||
else:
|
||||
assert messages[0]["reasoning"] == expected_reasoning
|
||||
if expected_tool_call_ids is None:
|
||||
assert "tool_calls" not in messages[0]
|
||||
else:
|
||||
assert [tool_call["id"] for tool_call in messages[0]["tool_calls"]] == (
|
||||
expected_tool_call_ids
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_returns_none_when_id_does_not_start_with_mcp_prefix(self):
|
||||
"""Test that non-MCP tool calls are not combined."""
|
||||
item = ResponseFunctionToolCall(
|
||||
type="function_call",
|
||||
id="regular_id", # Does not start with MCP_PREFIX
|
||||
call_id="call_123",
|
||||
name="test_function",
|
||||
arguments="{}",
|
||||
)
|
||||
messages = [{"role": "assistant", "reasoning": "some reasoning"}]
|
||||
|
||||
result = _maybe_combine_reasoning_and_tool_call(item, messages)
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_returns_none_when_last_message_is_not_assistant(self):
|
||||
"""Test that non-assistant last message returns None."""
|
||||
item = ResponseFunctionToolCall(
|
||||
type="function_call",
|
||||
id=f"{MCP_PREFIX}tool_id",
|
||||
call_id="call_123",
|
||||
name="test_function",
|
||||
arguments="{}",
|
||||
)
|
||||
messages = [{"role": "user", "content": "hello"}]
|
||||
|
||||
result = _maybe_combine_reasoning_and_tool_call(item, messages)
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_returns_none_when_last_message_has_no_reasoning(self):
|
||||
"""Test that assistant message without reasoning returns None."""
|
||||
item = ResponseFunctionToolCall(
|
||||
type="function_call",
|
||||
id=f"{MCP_PREFIX}tool_id",
|
||||
call_id="call_123",
|
||||
name="test_function",
|
||||
arguments="{}",
|
||||
)
|
||||
messages = [{"role": "assistant", "content": "some content"}]
|
||||
|
||||
result = _maybe_combine_reasoning_and_tool_call(item, messages)
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_combines_reasoning_and_mcp_tool_call(self):
|
||||
"""Test successful combination of reasoning message and MCP tool call."""
|
||||
item = ResponseFunctionToolCall(
|
||||
type="function_call",
|
||||
id=f"{MCP_PREFIX}tool_id",
|
||||
call_id="call_123",
|
||||
name="test_function",
|
||||
arguments='{"arg": "value"}',
|
||||
)
|
||||
messages = [{"role": "assistant", "reasoning": "I need to call this tool"}]
|
||||
|
||||
result = _maybe_combine_reasoning_and_tool_call(item, messages)
|
||||
|
||||
assert result is not None
|
||||
assert result["role"] == "assistant"
|
||||
assert result["reasoning"] == "I need to call this tool"
|
||||
assert "tool_calls" in result
|
||||
assert len(result["tool_calls"]) == 1
|
||||
assert result["tool_calls"][0]["id"] == "call_123"
|
||||
assert result["tool_calls"][0]["function"]["name"] == "test_function"
|
||||
assert result["tool_calls"][0]["function"]["arguments"] == '{"arg": "value"}'
|
||||
assert result["tool_calls"][0]["type"] == "function"
|
||||
|
||||
def test_returns_none_for_non_function_tool_call_type(self):
|
||||
"""Test that non-ResponseFunctionToolCall items return None."""
|
||||
# Pass a dict instead of ResponseFunctionToolCall
|
||||
item = {"type": "message", "content": "hello"}
|
||||
messages = [{"role": "assistant", "reasoning": "some reasoning"}]
|
||||
|
||||
result = _maybe_combine_reasoning_and_tool_call(item, messages)
|
||||
|
||||
assert result is None
|
||||
|
||||
def test_returns_none_when_id_is_empty_string(self):
|
||||
"""Test that empty string id returns None (falsy check)."""
|
||||
item = ResponseFunctionToolCall(
|
||||
type="function_call",
|
||||
id="", # Empty string is falsy
|
||||
call_id="call_123",
|
||||
name="test_function",
|
||||
arguments="{}",
|
||||
)
|
||||
messages = [{"role": "assistant", "reasoning": "some reasoning"}]
|
||||
|
||||
result = _maybe_combine_reasoning_and_tool_call(item, messages)
|
||||
|
||||
assert result is None
|
||||
@pytest.mark.parametrize(
|
||||
("items", "num_expected_messages"),
|
||||
[
|
||||
pytest.param(
|
||||
[
|
||||
make_output_message("Hello"),
|
||||
make_output_message("World"),
|
||||
],
|
||||
2,
|
||||
id="consecutive-output-messages",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
make_reasoning_item(content_text="Let me think"),
|
||||
make_reasoning_item(content_text="Let me think more"),
|
||||
],
|
||||
2,
|
||||
id="consecutive-reasoning-messages",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
make_function_call(call_id="call_123"),
|
||||
make_function_call_output(call_id="call_123", output="42"),
|
||||
make_function_call(call_id="call_456"),
|
||||
],
|
||||
3,
|
||||
id="interrupted-by-non-assistant-item",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_merge_chain_breaks(self, items, num_expected_messages):
|
||||
messages = construct_chat_messages_with_tool_call(items)
|
||||
assert len(messages) == num_expected_messages
|
||||
|
||||
|
||||
class TestConstructInputMessagesInstructionsLeak:
|
||||
|
||||
@@ -145,6 +145,7 @@ def test_openapi_stateless(case: Case):
|
||||
("POST", "/v1/chat/completions/batch"): LONG_TIMEOUT_SECONDS,
|
||||
("POST", "/v1/completions"): LONG_TIMEOUT_SECONDS,
|
||||
("POST", "/v1/messages"): LONG_TIMEOUT_SECONDS,
|
||||
("POST", "/inference/v1/generate"): LONG_TIMEOUT_SECONDS,
|
||||
}.get(key, DEFAULT_TIMEOUT_SECONDS)
|
||||
|
||||
# No need to verify SSL certificate for localhost
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for the disagg request/response protocol.
|
||||
|
||||
These tests intentionally avoid spinning up a server — they exercise the
|
||||
pydantic validators on ``GenerateRequest`` directly so they run fast and
|
||||
fail loudly if the validator semantics ever drift.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
from vllm.entrypoints.serve.disagg.protocol import GenerateRequest
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
|
||||
def _base_payload() -> dict:
|
||||
return {"token_ids": [1, 2, 3], "sampling_params": {}}
|
||||
|
||||
|
||||
def test_omitted_max_tokens_is_not_provided():
|
||||
"""Body without ``max_tokens`` must surface as 'not provided' so the
|
||||
server can apply its own default instead of the dataclass 16."""
|
||||
req = GenerateRequest.model_validate(_base_payload())
|
||||
# SamplingParams' dataclass default leaks through the parsed instance —
|
||||
# this is exactly the bug the server-side defaulting works around.
|
||||
assert req.sampling_params.max_tokens == 16
|
||||
assert not req.is_sampling_param_provided("max_tokens")
|
||||
|
||||
|
||||
def test_explicit_max_tokens_is_provided():
|
||||
"""Even when the client picks the same value as the dataclass default,
|
||||
it must register as explicitly set so the server won't override it."""
|
||||
payload = _base_payload()
|
||||
payload["sampling_params"] = {"max_tokens": 16}
|
||||
req = GenerateRequest.model_validate(payload)
|
||||
assert req.sampling_params.max_tokens == 16
|
||||
assert req.is_sampling_param_provided("max_tokens")
|
||||
|
||||
payload["sampling_params"] = {"max_tokens": 256}
|
||||
req = GenerateRequest.model_validate(payload)
|
||||
assert req.sampling_params.max_tokens == 256
|
||||
assert req.is_sampling_param_provided("max_tokens")
|
||||
|
||||
|
||||
def test_other_fields_tracked_independently():
|
||||
payload = _base_payload()
|
||||
payload["sampling_params"] = {"temperature": 0.0}
|
||||
req = GenerateRequest.model_validate(payload)
|
||||
assert not req.is_sampling_param_provided("max_tokens")
|
||||
assert req.is_sampling_param_provided("temperature")
|
||||
|
||||
|
||||
def test_json_roundtrip_preserves_provided_keys():
|
||||
payload = _base_payload()
|
||||
payload["sampling_params"] = {"temperature": 0.5}
|
||||
req = GenerateRequest.model_validate_json(json.dumps(payload))
|
||||
assert not req.is_sampling_param_provided("max_tokens")
|
||||
assert req.is_sampling_param_provided("temperature")
|
||||
|
||||
|
||||
def test_internal_instance_construction_treats_all_as_provided():
|
||||
"""When internal callers build ``GenerateRequest`` from a pre-resolved
|
||||
``SamplingParams`` instance, every field is considered explicitly set
|
||||
so server-side defaulting can't clobber values resolved upstream."""
|
||||
sp = SamplingParams(max_tokens=500, temperature=0.0)
|
||||
req = GenerateRequest(token_ids=[1, 2, 3], sampling_params=sp)
|
||||
assert req.is_sampling_param_provided("max_tokens")
|
||||
assert req.is_sampling_param_provided("temperature")
|
||||
# And keys we never touched should also count as provided in this path.
|
||||
assert req.is_sampling_param_provided("top_p")
|
||||
@@ -114,6 +114,36 @@ async def test_generate_endpoint(client):
|
||||
assert "choices" in data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_defaults_max_tokens_when_omitted(client):
|
||||
"""Regression: omitting ``max_tokens`` must not silently cap at 16.
|
||||
|
||||
``SamplingParams.max_tokens`` defaults to 16. Before the server-side
|
||||
defaulting was wired up, every request that didn't set ``max_tokens``
|
||||
truncated mid-generation. The server should now fill it in from
|
||||
``max_model_len - prompt_len`` (matching ``/v1/chat/completions``).
|
||||
"""
|
||||
payload = {
|
||||
"model": MODEL_NAME,
|
||||
"token_ids": [1, 2, 3],
|
||||
"sampling_params": {
|
||||
"temperature": 0.0,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
"stream": False,
|
||||
}
|
||||
resp = await client.post(GEN_ENDPOINT, json=payload)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
completion_tokens = len(data["choices"][0]["token_ids"])
|
||||
# max_model_len=1024 in the test fixture, prompt is 3 tokens, so we
|
||||
# should get ~1021 tokens of output (capped at max_model_len boundary).
|
||||
assert completion_tokens > 16, (
|
||||
f"expected server-side default to exceed the legacy 16-token cap, "
|
||||
f"got {completion_tokens}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generate_stream(client):
|
||||
payload = {
|
||||
|
||||
+214
-66
@@ -13,20 +13,40 @@ from torch.fx.experimental.proxy_tensor import make_fx
|
||||
import vllm.ir.op
|
||||
from vllm.ir.op import RESERVED_PROVIDERS, IrOp, IrOpImpl
|
||||
|
||||
# This should not exist
|
||||
assert "_custom_add" not in IrOp.registry
|
||||
|
||||
|
||||
class CustomError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@vllm.ir.register_op(allow_inplace=True)
|
||||
def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return x + y
|
||||
@pytest.fixture
|
||||
def custom_add_op(fake_vllm_ir):
|
||||
"""Register ``_custom_add`` plus impl_a, impl_b, impl_even for this test."""
|
||||
|
||||
@vllm.ir.register_op(allow_inplace=True)
|
||||
def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return x + y
|
||||
|
||||
@_custom_add.register_impl("impl_a")
|
||||
def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return x + y + 10
|
||||
|
||||
@_custom_add.register_impl("impl_b", inplace=True)
|
||||
def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
"""Computes x+y+20"""
|
||||
x.add_(y)
|
||||
x.add_(20)
|
||||
return x
|
||||
|
||||
@_custom_add.register_impl(
|
||||
"impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0
|
||||
)
|
||||
def impl_even(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return x + y + 50
|
||||
|
||||
return _custom_add
|
||||
|
||||
|
||||
def test_registration_overloads():
|
||||
def test_registration_overloads(fake_vllm_ir):
|
||||
assert all(
|
||||
n not in IrOp.registry for n in ["_custom_sub", "_custom_mul", "_custom_div"]
|
||||
)
|
||||
@@ -63,7 +83,7 @@ def test_registration_overloads():
|
||||
return x * y - 100
|
||||
|
||||
|
||||
def test_no_kw_only_args():
|
||||
def test_no_kw_only_args(fake_vllm_ir):
|
||||
# kw-only args not supported
|
||||
with pytest.raises(ValueError, match="keyword-only arguments"):
|
||||
|
||||
@@ -78,18 +98,23 @@ def test_no_kw_only_args():
|
||||
|
||||
class TestIrOpCustomAdd:
|
||||
# Registration invariants
|
||||
def test_decorated_object(self):
|
||||
def test_decorated_object(self, custom_add_op):
|
||||
"""Make sure that referring directly to an op is correct"""
|
||||
_custom_add = custom_add_op
|
||||
assert isinstance(_custom_add, IrOp)
|
||||
assert "_custom_add" in IrOp.registry
|
||||
assert _custom_add is IrOp.registry["_custom_add"]
|
||||
|
||||
def test_torch_op_is_registered(self):
|
||||
assert hasattr(torch.ops.vllm_ir, "_custom_add")
|
||||
assert callable(torch.ops.vllm_ir._custom_add.default)
|
||||
def test_torch_op_is_registered(self, custom_add_op):
|
||||
_custom_add = custom_add_op
|
||||
torch_ops = getattr(torch.ops, vllm.ir.op.vllm_ir_torch_lib.ns)
|
||||
assert hasattr(torch_ops, "_custom_add")
|
||||
assert callable(torch_ops._custom_add.default)
|
||||
assert _custom_add.torch_op is torch_ops._custom_add.default
|
||||
|
||||
# Semantic correctness
|
||||
def test_semantics_match_native(self):
|
||||
def test_semantics_match_native(self, custom_add_op):
|
||||
_custom_add = custom_add_op
|
||||
x = torch.randn(4, 5)
|
||||
y = torch.randn(4, 5)
|
||||
|
||||
@@ -103,7 +128,9 @@ class TestIrOpCustomAdd:
|
||||
# Implementation registration
|
||||
# -------------------------
|
||||
|
||||
def test_register_impl_is_non_intrusive(self):
|
||||
def test_register_impl_is_non_intrusive(self, custom_add_op):
|
||||
_custom_add = custom_add_op
|
||||
|
||||
@_custom_add.register_impl("dummy_provider")
|
||||
def dummy_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return x + y + 123
|
||||
@@ -117,7 +144,8 @@ class TestIrOpCustomAdd:
|
||||
# Native semantics must still hold
|
||||
torch.testing.assert_close(_custom_add(x, y), x + y)
|
||||
|
||||
def test_schema_contains_tensor_signature(self):
|
||||
def test_schema_contains_tensor_signature(self, custom_add_op):
|
||||
_custom_add = custom_add_op
|
||||
schema = _custom_add._schema_str
|
||||
|
||||
assert "Tensor" in schema
|
||||
@@ -131,10 +159,19 @@ class TestIrOpCustomAdd:
|
||||
@pytest.mark.parametrize("symbolic_trace", [True, False])
|
||||
@pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
|
||||
def test_trace_sees_single_custom_op(
|
||||
self, symbolic_trace: bool, enable_torch_wrap: bool, overload: str
|
||||
self,
|
||||
custom_add_op,
|
||||
symbolic_trace: bool,
|
||||
enable_torch_wrap: bool,
|
||||
overload: str,
|
||||
):
|
||||
_custom_add = custom_add_op
|
||||
op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace
|
||||
torch_op = getattr(torch.ops.vllm_ir._custom_add, overload)
|
||||
torch_op = (
|
||||
_custom_add.torch_op
|
||||
if overload == "default"
|
||||
else _custom_add.maybe_inplace.torch_op
|
||||
)
|
||||
|
||||
def fn(x, y):
|
||||
return op_fn(x, y)
|
||||
@@ -175,30 +212,13 @@ class TestIrOpCustomAdd:
|
||||
assert len(ir_nodes) == 1, gm.code
|
||||
|
||||
|
||||
@_custom_add.register_impl("impl_a")
|
||||
def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return x + y + 10
|
||||
|
||||
|
||||
@_custom_add.register_impl("impl_b", inplace=True)
|
||||
def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
"""Computes x+y+20"""
|
||||
x.add_(y)
|
||||
x.add_(20)
|
||||
return x
|
||||
|
||||
|
||||
@_custom_add.register_impl("impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0)
|
||||
def impl_even(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return x + y + 50
|
||||
|
||||
|
||||
class TestIrOpImplDispatch:
|
||||
def test_register_impl(self):
|
||||
def test_register_impl(self, custom_add_op):
|
||||
_custom_add = custom_add_op
|
||||
assert "impl_a" in _custom_add.impls
|
||||
impl = _custom_add.impls["impl_a"]
|
||||
|
||||
assert impl is impl_a
|
||||
assert impl is _custom_add.impls["impl_a"]
|
||||
assert impl.op is _custom_add
|
||||
assert impl.provider == "impl_a"
|
||||
assert callable(impl.impl_fn)
|
||||
@@ -211,14 +231,15 @@ class TestIrOpImplDispatch:
|
||||
return x + y + 30
|
||||
|
||||
# Check the original impl is still intact
|
||||
assert _custom_add.impls["impl_a"] is impl_a
|
||||
assert _custom_add.impls["impl_a"] is impl
|
||||
|
||||
# Check support all args
|
||||
assert impl_a.supports_all_args
|
||||
assert impl_b.supports_all_args
|
||||
assert not impl_even.supports_all_args
|
||||
assert _custom_add.impls["impl_a"].supports_all_args
|
||||
assert _custom_add.impls["impl_b"].supports_all_args
|
||||
assert not _custom_add.impls["impl_even"].supports_all_args
|
||||
|
||||
def test_reserved_provider_rejected(self):
|
||||
def test_reserved_provider_rejected(self, custom_add_op):
|
||||
_custom_add = custom_add_op
|
||||
for provider in RESERVED_PROVIDERS:
|
||||
with pytest.raises(AssertionError):
|
||||
|
||||
@@ -226,7 +247,8 @@ class TestIrOpImplDispatch:
|
||||
def bad_impl(x, y):
|
||||
return x + y
|
||||
|
||||
def test_set_priority_scoped(self):
|
||||
def test_set_priority_scoped(self, custom_add_op):
|
||||
_custom_add = custom_add_op
|
||||
assert _custom_add.get_priority() == []
|
||||
|
||||
with _custom_add.set_priority(["impl_even", "impl_b"]):
|
||||
@@ -251,20 +273,25 @@ class TestIrOpImplDispatch:
|
||||
assert _custom_add.get_priority() == []
|
||||
|
||||
@pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
|
||||
def test_dispatch_priority_order(self, overload: str):
|
||||
def test_dispatch_priority_order(self, custom_add_op, overload: str):
|
||||
_custom_add = custom_add_op
|
||||
op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace
|
||||
torch_op = getattr(torch.ops.vllm_ir._custom_add, overload)
|
||||
torch_op = (
|
||||
_custom_add.torch_op
|
||||
if overload == "default"
|
||||
else _custom_add.maybe_inplace.torch_op
|
||||
)
|
||||
|
||||
x = torch.tensor(1, dtype=torch.int32)
|
||||
y = torch.tensor(2, dtype=torch.int32)
|
||||
|
||||
with _custom_add.set_priority(["impl_b", "impl_a"]):
|
||||
assert _custom_add.dispatch(x, y) is impl_b
|
||||
assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_b"]
|
||||
out1 = op_fn(x.clone(), y)
|
||||
out2 = torch_op(x.clone(), y)
|
||||
|
||||
with _custom_add.set_priority(["impl_a"]):
|
||||
assert _custom_add.dispatch(x, y) is impl_a
|
||||
assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_a"]
|
||||
out3 = op_fn(x.clone(), y)
|
||||
out4 = torch_op(x.clone(), y)
|
||||
|
||||
@@ -275,7 +302,9 @@ class TestIrOpImplDispatch:
|
||||
assert out3.item() == 1 + 2 + 10
|
||||
assert out4.item() == 1 + 2 + 10
|
||||
|
||||
def test_unsupported_impl_filtered(self):
|
||||
def test_unsupported_impl_filtered(self, custom_add_op):
|
||||
_custom_add = custom_add_op
|
||||
|
||||
@_custom_add.register_impl("impl_unsupported", supported=False)
|
||||
def impl_unsupported(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
return x + y + 999
|
||||
@@ -291,8 +320,9 @@ class TestIrOpImplDispatch:
|
||||
assert out.item() == 1 + 2 + 10
|
||||
|
||||
def test_supports_args_runtime_dispatch_and_warning(
|
||||
self, caplog_vllm: pytest.LogCaptureFixture
|
||||
self, custom_add_op, caplog_vllm: pytest.LogCaptureFixture
|
||||
):
|
||||
_custom_add = custom_add_op
|
||||
x1 = torch.ones((2, 2), dtype=torch.int32)
|
||||
y1 = torch.full((2, 2), 2, dtype=torch.int32)
|
||||
|
||||
@@ -312,7 +342,7 @@ class TestIrOpImplDispatch:
|
||||
|
||||
# Check dispatching
|
||||
assert _custom_add.get_priority() == ["impl_even", "native"]
|
||||
assert _custom_add.dispatch(x1, y1) is impl_even
|
||||
assert _custom_add.dispatch(x1, y1) is _custom_add.impls["impl_even"]
|
||||
assert _custom_add.dispatch(x2, y2) is _custom_add.impls["native"]
|
||||
|
||||
out1 = _custom_add(x1, y1) # size(1) == 2 → impl_even
|
||||
@@ -324,8 +354,12 @@ class TestIrOpImplDispatch:
|
||||
assert torch.all(out2 == 1 + 2)
|
||||
|
||||
def test_default_priority(
|
||||
self, caplog_vllm: pytest.LogCaptureFixture, disable_log_dedup
|
||||
self,
|
||||
custom_add_op,
|
||||
caplog_vllm: pytest.LogCaptureFixture,
|
||||
disable_log_dedup,
|
||||
):
|
||||
_custom_add = custom_add_op
|
||||
# Make sure logs are not deduplicated to properly test the warning
|
||||
x = torch.tensor([3], dtype=torch.int32)
|
||||
y = torch.tensor([4], dtype=torch.int32)
|
||||
@@ -347,15 +381,23 @@ class TestIrOpImplDispatch:
|
||||
assert "priority not set" in message
|
||||
|
||||
|
||||
@vllm.ir.register_op
|
||||
def _custom_mm(
|
||||
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
|
||||
) -> torch.Tensor:
|
||||
tmp = x @ y
|
||||
return tmp if bias is None else tmp + bias
|
||||
@pytest.fixture
|
||||
def custom_mm_op(fake_vllm_ir):
|
||||
"""Fixture that registers ``_custom_mm`` (isolated by ``fake_vllm_ir``)."""
|
||||
|
||||
@vllm.ir.register_op
|
||||
def _custom_mm(
|
||||
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
|
||||
) -> torch.Tensor:
|
||||
tmp = x @ y
|
||||
return tmp if bias is None else tmp + bias
|
||||
|
||||
return _custom_mm
|
||||
|
||||
|
||||
def test_default_args():
|
||||
def test_default_args(custom_mm_op):
|
||||
_custom_mm = custom_mm_op
|
||||
|
||||
# Test that default args are properly applied when dispatching and calling
|
||||
@_custom_mm.register_impl("impl_mm", supports_args=lambda x, y, bias=None: True)
|
||||
def impl_mm(
|
||||
@@ -373,7 +415,8 @@ def test_default_args():
|
||||
assert _custom_mm.dispatch(x1, x2) is impl_mm
|
||||
|
||||
|
||||
def test_bad_impl_registrations():
|
||||
def test_bad_impl_registrations(custom_mm_op):
|
||||
_custom_mm = custom_mm_op
|
||||
# Check bad schema
|
||||
with pytest.raises(ValueError, match="does not match native schema"):
|
||||
|
||||
@@ -446,7 +489,8 @@ def test_bad_impl_registrations():
|
||||
) -> torch.Tensor:
|
||||
return x @ y + 40
|
||||
|
||||
assert set(_custom_mm.impls.keys()) == {"impl_mm", "native"}
|
||||
# With fixture, each test gets a fresh op with only "native" impl
|
||||
assert set(_custom_mm.impls.keys()) == {"native"}
|
||||
|
||||
|
||||
IMPL_OOT_SRC = """
|
||||
@@ -460,14 +504,14 @@ def impl_mm_oot(
|
||||
"""
|
||||
|
||||
|
||||
def load_custom_mm_module(file_path: Path):
|
||||
def load_custom_mm_module(file_path: Path, custom_mm_op):
|
||||
spec = importlib.util.spec_from_file_location("_custom_mm_oot", file_path)
|
||||
assert spec is not None
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
|
||||
# Inject the variable into the module's global namespace
|
||||
# This allows the @_custom_mm.register_impl decorator to work
|
||||
module._custom_mm = _custom_mm # type: ignore[attr-defined]
|
||||
module._custom_mm = custom_mm_op # type: ignore[attr-defined]
|
||||
|
||||
# Execute the file; this triggers the decorator
|
||||
assert spec.loader is not None
|
||||
@@ -475,12 +519,13 @@ def load_custom_mm_module(file_path: Path):
|
||||
return module
|
||||
|
||||
|
||||
def test_uuid_and_oot(tmp_path: Path):
|
||||
def test_uuid_and_oot(custom_mm_op, tmp_path: Path):
|
||||
_custom_mm = custom_mm_op
|
||||
file_path = tmp_path / "_custom_mm_oot.py"
|
||||
file_path.write_text(IMPL_OOT_SRC)
|
||||
|
||||
assert "impl_mm_oot" not in _custom_mm.impls
|
||||
_ = load_custom_mm_module(file_path)
|
||||
_ = load_custom_mm_module(file_path, _custom_mm)
|
||||
assert "impl_mm_oot" in _custom_mm.impls
|
||||
|
||||
uuid = _custom_mm.impls["impl_mm_oot"].uuid()
|
||||
@@ -489,7 +534,7 @@ def test_uuid_and_oot(tmp_path: Path):
|
||||
# Replace file source
|
||||
file_path.write_text(IMPL_OOT_SRC + " # added file source")
|
||||
assert "impl_mm_oot" not in _custom_mm.impls
|
||||
_ = load_custom_mm_module(file_path)
|
||||
_ = load_custom_mm_module(file_path, _custom_mm)
|
||||
assert "impl_mm_oot" in _custom_mm.impls
|
||||
|
||||
uuid1 = _custom_mm.impls["impl_mm_oot"].uuid()
|
||||
@@ -499,7 +544,7 @@ def test_uuid_and_oot(tmp_path: Path):
|
||||
# Back to original
|
||||
file_path.write_text(IMPL_OOT_SRC)
|
||||
assert "impl_mm_oot" not in _custom_mm.impls
|
||||
_ = load_custom_mm_module(file_path)
|
||||
_ = load_custom_mm_module(file_path, _custom_mm)
|
||||
assert "impl_mm_oot" in _custom_mm.impls
|
||||
|
||||
uuid2 = _custom_mm.impls["impl_mm_oot"].uuid()
|
||||
@@ -571,3 +616,106 @@ class TestTolerance:
|
||||
op = IrOp("_tol_test_unknown", _test_native)
|
||||
with pytest.raises(ValueError, match="No tolerance defined"):
|
||||
op.get_tolerance(torch.complex64)
|
||||
|
||||
|
||||
def test_naming_validation(fake_vllm_ir):
|
||||
"""Test that op and provider names are validated ([a-z_][a-z_0-9]*)."""
|
||||
|
||||
# Valid op and provider names
|
||||
@vllm.ir.register_op
|
||||
def _valid_name_123(x: torch.Tensor) -> torch.Tensor:
|
||||
return x
|
||||
|
||||
@_valid_name_123.register_impl("valid_provider_123")
|
||||
def valid_impl(x: torch.Tensor) -> torch.Tensor:
|
||||
return x + 1
|
||||
|
||||
# Invalid op names should fail
|
||||
with pytest.raises(ValueError, match="name.*invalid"):
|
||||
|
||||
@vllm.ir.register_op
|
||||
def InvalidName(x: torch.Tensor) -> torch.Tensor:
|
||||
return x
|
||||
|
||||
with pytest.raises(ValueError, match="name.*invalid"):
|
||||
|
||||
@vllm.ir.register_op(name="123invalid")
|
||||
def some_func(x: torch.Tensor) -> torch.Tensor:
|
||||
return x
|
||||
|
||||
# Invalid provider names should fail
|
||||
with pytest.raises(ValueError, match="name.*invalid"):
|
||||
|
||||
@_valid_name_123.register_impl("Invalid-Provider")
|
||||
def invalid_impl(x: torch.Tensor) -> torch.Tensor:
|
||||
return x + 1
|
||||
|
||||
|
||||
def test_registration_stack_traces(fake_vllm_ir):
|
||||
"""Test that stack traces are captured for ops and impls."""
|
||||
|
||||
@vllm.ir.register_op
|
||||
def _test_stack(x: torch.Tensor) -> torch.Tensor:
|
||||
return x
|
||||
|
||||
@_test_stack.register_impl("test_provider")
|
||||
def test_impl(x: torch.Tensor) -> torch.Tensor:
|
||||
return x + 1
|
||||
|
||||
# Verify op stack trace
|
||||
assert hasattr(_test_stack, "_registration_stack")
|
||||
assert len(_test_stack._registration_stack) > 0
|
||||
op_stack_str = "".join(_test_stack._registration_stack)
|
||||
assert "test_op.py" in op_stack_str
|
||||
# Last frame should be the decorator in user code, not internal decorator logic
|
||||
assert "@vllm.ir.register_op" in _test_stack._registration_stack[-1]
|
||||
assert "return decorator(f)" not in op_stack_str
|
||||
|
||||
# Verify impl stack trace
|
||||
impl = _test_stack.impls["test_provider"]
|
||||
assert hasattr(impl, "_registration_stack")
|
||||
assert len(impl._registration_stack) > 0
|
||||
impl_stack_str = "".join(impl._registration_stack)
|
||||
assert "test_op.py" in impl_stack_str
|
||||
# Last frame should be the decorator in user code
|
||||
assert '@_test_stack.register_impl("test_provider")' in impl._registration_stack[-1]
|
||||
|
||||
|
||||
def test_op_repr_uses_docstring(fake_vllm_ir):
|
||||
"""Test that __str__ uses the function's docstring and __repr__ is simple."""
|
||||
|
||||
@vllm.ir.register_op
|
||||
def _test_repr_with_doc(x: torch.Tensor) -> torch.Tensor:
|
||||
"""First line of docstring.
|
||||
|
||||
Additional details here.
|
||||
"""
|
||||
return x
|
||||
|
||||
@vllm.ir.register_op
|
||||
def _test_repr_no_doc(x: torch.Tensor) -> torch.Tensor:
|
||||
return x
|
||||
|
||||
# __str__ with docstring: uses first line only
|
||||
str_with = str(_test_repr_with_doc)
|
||||
assert "IrOp('_test_repr_with_doc')" in str_with
|
||||
assert "First line of docstring." in str_with
|
||||
assert "Additional details" not in str_with
|
||||
|
||||
# __str__ without docstring: simple format
|
||||
assert str(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')"
|
||||
|
||||
# __repr__ should be simple for both
|
||||
assert repr(_test_repr_with_doc) == "IrOp('_test_repr_with_doc')"
|
||||
assert repr(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')"
|
||||
|
||||
|
||||
def test_vllm_ir_fixture(fake_vllm_ir):
|
||||
"""Test that the fake_vllm_ir fixture provides test isolation."""
|
||||
|
||||
@vllm.ir.register_op
|
||||
def _test_fixture(x: torch.Tensor) -> torch.Tensor:
|
||||
return x
|
||||
|
||||
assert "_test_fixture" in IrOp.registry
|
||||
# Fixture will automatically clean up after test
|
||||
|
||||
@@ -345,3 +345,182 @@ def test_triton_unified_attn_fp16_input_fp8_output(
|
||||
torch.testing.assert_close(output_fp16, ref_output, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output_fp16 - ref_output))}",
|
||||
)
|
||||
|
||||
|
||||
# USE_TD path covers two head-size regimes:
|
||||
# - pow2 (HEAD_SIZE == HEAD_SIZE_PADDED): full TD path including Q/O.
|
||||
# - non-pow2 (96, HEAD_SIZE_PADDED=128): gates USE_TD_QO off — Q load
|
||||
# and output store fall back to pointer path, KV tile TD load remains.
|
||||
# The non-pow2 case mirrors real models like Phi-3-mini (head_size=96).
|
||||
HEAD_SIZES_USE_TD = [128, 256, 96]
|
||||
|
||||
|
||||
def _run_use_td_case(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
sliding_window: int | None,
|
||||
soft_cap: float | None,
|
||||
seq_threshold_3D: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
num_blocks: int = 2048,
|
||||
) -> None:
|
||||
"""Shared driver for the USE_TD test cases.
|
||||
|
||||
Runs ``unified_attention(..., use_td=True)`` and compares against the
|
||||
reference paged-attention implementation that the sibling non-TD
|
||||
tests use.
|
||||
"""
|
||||
torch.set_default_device(DEVICE_TYPE)
|
||||
set_random_seed(0)
|
||||
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads, num_kv_heads = num_heads
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
output = torch.empty_like(query)
|
||||
|
||||
num_par_softmax_segments = 16
|
||||
head_size_padded = next_power_of_2(head_size)
|
||||
softmax_segm_output = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_max = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_expsum = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
unified_attention(
|
||||
q=query,
|
||||
k=key_cache,
|
||||
v=value_cache,
|
||||
out=output,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
seqused_k=kv_lens_tensor,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
block_table=block_tables,
|
||||
softcap=soft_cap if soft_cap is not None else 0,
|
||||
q_descale=None,
|
||||
k_descale=None,
|
||||
v_descale=None,
|
||||
seq_threshold_3D=seq_threshold_3D,
|
||||
num_par_softmax_segments=num_par_softmax_segments,
|
||||
softmax_segm_output=softmax_segm_output,
|
||||
softmax_segm_max=softmax_segm_max,
|
||||
softmax_segm_expsum=softmax_segm_expsum,
|
||||
use_td=True,
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
sliding_window=sliding_window,
|
||||
soft_cap=soft_cap,
|
||||
)
|
||||
torch.testing.assert_close(output, ref_output, atol=1.5e-2, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens", [[(1, 1328), (5, 18), (129, 463)], [(1, 523), (1, 37), (1, 2011)]]
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES_USE_TD)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("sliding_window", [None, 128])
|
||||
@pytest.mark.parametrize("soft_cap", [None, 50.0])
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
|
||||
@torch.inference_mode()
|
||||
def test_triton_unified_attn_use_td(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int | None,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
seq_threshold_3D: int,
|
||||
) -> None:
|
||||
"""Exercise the USE_TD (tensor-descriptor) Q/K/V load/store path.
|
||||
|
||||
Covers both 2D and 3D kernels via ``seq_threshold_3D``. Two routes
|
||||
to the USE_TD_QO=False fallback (pointer path for Q/O with TD still
|
||||
active for KV tile loads):
|
||||
|
||||
- non-pow2 ``num_queries_per_kv`` via ``NUM_HEADS`` entry ``(5, 1)``,
|
||||
- non-pow2 ``head_size`` via ``HEAD_SIZES_USE_TD`` entry ``96``.
|
||||
"""
|
||||
_run_use_td_case(
|
||||
seq_lens=seq_lens,
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
block_size=block_size,
|
||||
sliding_window=sliding_window,
|
||||
soft_cap=soft_cap,
|
||||
seq_threshold_3D=seq_threshold_3D,
|
||||
num_blocks=num_blocks,
|
||||
)
|
||||
|
||||
|
||||
# Prefill-heavy shape: long query drives the prefill kernel path where
|
||||
# ``_get_tile_size`` returns 32, which exceeds block_size=16 and must be
|
||||
# clamped by the fix in 'clamp TILE_SIZE to block_size when USE_TD'.
|
||||
# Only the prefill launch exercises the clamp, so parameterize only over
|
||||
# the (num_heads, seq_threshold_3D=0) combinations needed to cover it.
|
||||
@pytest.mark.parametrize("num_heads", [(4, 4), (5, 1)])
|
||||
@torch.inference_mode()
|
||||
def test_triton_unified_attn_use_td_tile_clamp(
|
||||
num_heads: tuple[int, int],
|
||||
) -> None:
|
||||
"""Regression guard: ``USE_TD`` needs ``BLOCK_SIZE % TILE_SIZE == 0``.
|
||||
|
||||
With ``block_size=16`` and ``head_size=128`` (non-Gemma3),
|
||||
``_get_tile_size`` returns 32 for prefill, which violates the
|
||||
``USE_TD`` constraint unless clamped to ``block_size``. Without
|
||||
the clamp the triton kernel ``static_assert`` fires at compile time.
|
||||
"""
|
||||
_run_use_td_case(
|
||||
seq_lens=[(256, 256), (128, 128)],
|
||||
num_heads=num_heads,
|
||||
head_size=128,
|
||||
block_size=16,
|
||||
sliding_window=None,
|
||||
soft_cap=None,
|
||||
seq_threshold_3D=0,
|
||||
)
|
||||
|
||||
@@ -15,16 +15,16 @@ from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import (
|
||||
BatchedDeepGemmExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import DeepGemmExperts
|
||||
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
|
||||
from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
|
||||
BatchedTritonExperts,
|
||||
NaiveBatchedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
|
||||
TritonOrDeepGemmExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
|
||||
MoEPrepareAndFinalizeNoDPEPModular,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
|
||||
TritonOrDeepGemmExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
|
||||
cutlass_fp4_supported,
|
||||
)
|
||||
|
||||
@@ -9,7 +9,7 @@ from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_confi
|
||||
from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import (
|
||||
BatchedDeepGemmExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
|
||||
from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
|
||||
BatchedTritonExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
|
||||
|
||||
@@ -16,7 +16,7 @@ from tests.kernels.quant_utils import native_batched_masked_quant_matmul
|
||||
from tests.kernels.utils import torch_experts
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.fused_moe import fused_topk
|
||||
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
|
||||
from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
|
||||
invoke_moe_batched_triton_kernel,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
@@ -31,7 +31,7 @@ from vllm.model_executor.layers.fused_moe.config import (
|
||||
from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
|
||||
_valid_deep_gemm_shape,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
|
||||
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
|
||||
TritonOrDeepGemmExperts,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
@@ -19,7 +19,9 @@ from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_batched_moe import BatchedTritonExperts
|
||||
from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
|
||||
BatchedTritonExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
per_token_group_quant_fp8,
|
||||
|
||||
@@ -25,10 +25,10 @@ from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantDesc,
|
||||
fp8_w8a8_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
|
||||
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
|
||||
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
|
||||
TritonOrDeepGemmExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
per_token_group_quant_fp8,
|
||||
)
|
||||
|
||||
@@ -1588,7 +1588,7 @@ def test_unquantized_bf16_flashinfer_trtllm_backend(
|
||||
layer.apply_router_weight_on_input = False
|
||||
layer.routed_scaling_factor = None
|
||||
layer.shared_experts = None
|
||||
layer._maybe_init_expert_routing_tables = lambda: None
|
||||
layer._expert_routing_tables = lambda: None
|
||||
|
||||
quant_method.process_weights_after_loading(layer)
|
||||
|
||||
|
||||
@@ -31,12 +31,13 @@ from vllm.config import (
|
||||
VllmConfig,
|
||||
set_current_vllm_config,
|
||||
)
|
||||
from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
|
||||
from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
|
||||
from vllm.distributed.parallel_state import (
|
||||
from vllm.distributed import (
|
||||
get_ep_group,
|
||||
get_eplb_group,
|
||||
tensor_model_parallel_all_gather,
|
||||
)
|
||||
from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
|
||||
from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
|
||||
from vllm.forward_context import set_forward_context
|
||||
from vllm.model_executor.layers.fused_moe import FusedMoE, fused_experts
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
@@ -49,6 +50,7 @@ from vllm.model_executor.layers.quantization.modelopt import (
|
||||
ModelOptFp8Config,
|
||||
ModelOptNvFp4Config,
|
||||
)
|
||||
from vllm.model_executor.models.utils import sequence_parallel_chunk
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.flashinfer import (
|
||||
has_flashinfer_nvlink_one_sided,
|
||||
@@ -81,6 +83,9 @@ PARALLEL_COMBOS = [
|
||||
[1, 4, False],
|
||||
[2, 1, True],
|
||||
[4, 1, True],
|
||||
# This combination indicates sequence parallel.
|
||||
# See ParallelConfig.use_sequence_parallel.
|
||||
[2, 2, True],
|
||||
]
|
||||
|
||||
# TODO: should this even be set manually? let oracles handle this
|
||||
@@ -112,24 +117,24 @@ QUANT_METHODS = [
|
||||
# Which quantization methods each backend supports.
|
||||
# fmt: off
|
||||
BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {
|
||||
"allgather_reducescatter": {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
|
||||
"mori": {None, "fp8", "modelopt_fp8"},
|
||||
"flashinfer_nvlink_two_sided": {None, "modelopt_fp8", "modelopt_fp4"},
|
||||
"flashinfer_nvlink_one_sided": {None, "modelopt_fp8", "modelopt_fp4"},
|
||||
"deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"},
|
||||
"allgather_reducescatter": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501
|
||||
"mori": {None, "fp8", "modelopt_fp8"},
|
||||
"flashinfer_nvlink_two_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501
|
||||
"flashinfer_nvlink_one_sided": {None, "modelopt_fp4"}, # noqa: E501
|
||||
"deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501
|
||||
"deepep_high_throughput": {None, "fp8_blocked", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501
|
||||
"nixl_ep": {None, "fp8", "modelopt_fp8"},
|
||||
"nixl_ep": {None, "fp8_blocked", "modelopt_fp8"},
|
||||
}
|
||||
|
||||
# Map from backend -> (DP/EP support, DP support, TP support)
|
||||
BACKEND_EP_DP_TP_SUPPORT: dict[str, tuple[bool, bool, bool]] = {
|
||||
"allgather_reducescatter": (True, True, True),
|
||||
"mori": (True, False, False),
|
||||
"flashinfer_nvlink_two_sided": (False, True, False),
|
||||
"flashinfer_nvlink_one_sided": (False, True, False),
|
||||
"deepep_low_latency": (True, False, False),
|
||||
"deepep_high_throughput": (True, False, False),
|
||||
"nixl_ep": (True, False, False),
|
||||
# Map from backend -> (DP/EP support, DP support, TP support, SP support)
|
||||
BACKEND_EP_DP_TP_SUPPORT: dict[str, tuple[bool, bool, bool, bool]] = {
|
||||
"allgather_reducescatter": (True, True, True, True),
|
||||
"mori": (True, False, False, True),
|
||||
"flashinfer_nvlink_two_sided": (False, True, False, False),
|
||||
"flashinfer_nvlink_one_sided": (False, True, False, False),
|
||||
"deepep_low_latency": (True, False, False, True),
|
||||
"deepep_high_throughput": (True, False, False, True),
|
||||
"nixl_ep": (True, False, False, True),
|
||||
}
|
||||
# fmt: on
|
||||
|
||||
@@ -163,6 +168,45 @@ def override_normalize_e4m3fn_to_e4m3fnuz():
|
||||
vllm.model_executor.layers.quantization.utils.w8a8_utils.normalize_e4m3fn_to_e4m3fnuz = mock_normalize_e4m3fn_to_e4m3fnuz # noqa: E501
|
||||
|
||||
|
||||
def sp_wrapper(
|
||||
fn: Callable | FusedMoE, is_sequence_parallel: bool | None = None
|
||||
) -> Callable:
|
||||
"""Wrapper to handle sequence parallelism chunking and gathering.
|
||||
|
||||
For SP with EP:
|
||||
- The TP group is created with the original tensor_parallel_size (e.g., 2)
|
||||
- get_tp_group() has the correct world_size for SP operations
|
||||
- sequence_parallel_chunk() uses get_tensor_model_parallel_world_size()
|
||||
- tensor_model_parallel_all_gather() uses get_tp_group()
|
||||
- Both should work correctly even when EP is enabled
|
||||
"""
|
||||
if isinstance(fn, FusedMoE):
|
||||
assert is_sequence_parallel is None
|
||||
is_sequence_parallel = fn.is_sequence_parallel
|
||||
else:
|
||||
assert is_sequence_parallel is not None
|
||||
|
||||
if is_sequence_parallel:
|
||||
|
||||
def wrapper(
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
# Split sequence across TP ranks
|
||||
# Both hidden_states and router_logits have [num_tokens, ...] shape
|
||||
hidden_states = sequence_parallel_chunk(hidden_states)
|
||||
router_logits = sequence_parallel_chunk(router_logits)
|
||||
# Run MoE on local chunk
|
||||
result = fn(hidden_states, router_logits)
|
||||
# Gather results from all TP ranks
|
||||
result = tensor_model_parallel_all_gather(result, 0)
|
||||
# Remove any padding added by SP.
|
||||
return result[: hidden_states.shape[0]]
|
||||
|
||||
return wrapper
|
||||
return fn
|
||||
|
||||
|
||||
def maybe_roundup_layer_hidden_size(
|
||||
hidden_size: int,
|
||||
act_dtype: torch.dtype,
|
||||
@@ -272,6 +316,15 @@ class MoETestConfig:
|
||||
dp_size: int = 1
|
||||
tp_size: int = 1
|
||||
|
||||
@property
|
||||
def is_sequence_parallel(self) -> bool:
|
||||
# Sequence parallelism: EP enabled + TP dimension used for sequence splitting
|
||||
# In test config: ep_size represents total expert parallel size
|
||||
# tp_size represents the original TP dimension (becomes sp_size in FusedMoE)
|
||||
# dp_size represents data parallel size
|
||||
# For SP: we need EP enabled (ep_size > 1) and sequence splitting (tp_size > 1)
|
||||
return self.ep_size > 1 and self.tp_size > 1
|
||||
|
||||
# TODO: add more error messages
|
||||
def id(self) -> str:
|
||||
def proc(s: str) -> str:
|
||||
@@ -404,11 +457,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
|
||||
"leads to large differences.",
|
||||
)
|
||||
|
||||
# gate requires shared_experts (use_overlapped mode)
|
||||
# TODO: also not sure this is true
|
||||
if config.use_gate and not config.use_shared_experts:
|
||||
return False, "gate requires shared_experts (use_overlapped mode)"
|
||||
|
||||
# Skip modelopt_fp4 if not on B100+ (compute capability 10.0+)
|
||||
if (
|
||||
config.quantization == "modelopt_fp4"
|
||||
@@ -445,7 +493,7 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
|
||||
)
|
||||
|
||||
if config.backend == "nixl_ep":
|
||||
from vllm.model_executor.layers.fused_moe.nixl_ep_prepare_finalize import ( # noqa: E501
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.nixl_ep import ( # noqa: E501
|
||||
NixlEPPrepareAndFinalize,
|
||||
)
|
||||
|
||||
@@ -456,11 +504,11 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
|
||||
)
|
||||
|
||||
if config.backend is not None:
|
||||
supports_ep_dp, supports_dp, supports_tp = BACKEND_EP_DP_TP_SUPPORT[
|
||||
config.backend
|
||||
]
|
||||
supports_ep_dp, supports_dp, supports_tp, supports_sp = (
|
||||
BACKEND_EP_DP_TP_SUPPORT[config.backend]
|
||||
)
|
||||
|
||||
if config.tp_size > 1 and not supports_tp:
|
||||
if config.tp_size > 1 and not supports_tp and not config.is_sequence_parallel:
|
||||
return False, f"{config.backend} does not support TP."
|
||||
|
||||
if config.dp_size > 1 and config.ep_size == 1 and not supports_dp:
|
||||
@@ -468,10 +516,34 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
|
||||
|
||||
if config.dp_size > 1 and config.ep_size > 1 and not supports_ep_dp:
|
||||
return False, f"{config.backend} does not support EP/DP."
|
||||
|
||||
if config.is_sequence_parallel and not supports_sp:
|
||||
return False, f"{config.backend} does not support SP."
|
||||
else:
|
||||
if config.tp_size > 1 or config.ep_size > 1 or config.dp_size > 1:
|
||||
return False, "An all2all backend is required for parallelism."
|
||||
|
||||
# Sequence parallelism specific validations
|
||||
if config.is_sequence_parallel:
|
||||
if config.ep_size == 1:
|
||||
return False, "Sequence parallelism requires EP to be enabled (ep_size > 1)"
|
||||
|
||||
if config.tp_size == 1:
|
||||
return (
|
||||
False,
|
||||
"Sequence parallelism requires tp_size > 1 for sequence splitting",
|
||||
)
|
||||
|
||||
# SP is essentially EP + sequence splitting
|
||||
# Verify the relationship: ep_size should equal dp_size * tp_size
|
||||
# (when pcp_size=1).
|
||||
expected_ep_size = config.dp_size * config.tp_size
|
||||
if config.ep_size != expected_ep_size:
|
||||
return False, (
|
||||
f"For sequence parallelism: ep_size ({config.ep_size}) should equal "
|
||||
f"dp_size * tp_size ({expected_ep_size})"
|
||||
)
|
||||
|
||||
if config.enable_eplb:
|
||||
if config.ep_size == 1:
|
||||
return False, "EPLB requires EP."
|
||||
@@ -485,14 +557,6 @@ def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
|
||||
if config.num_experts % config.dp_size != 0:
|
||||
return False, "EPLB requires num_experts divisible by ep_size"
|
||||
|
||||
# Disable fp4 tests until flashinfer is updated or the Dockerfile is
|
||||
# modified to install cublasLt.h. See #39525.
|
||||
if (
|
||||
config.quantization == "modelopt_fp4"
|
||||
and current_platform.is_device_capability_family(100)
|
||||
):
|
||||
return False, "Temporarily skip until #39525 is resolved"
|
||||
|
||||
return True, None
|
||||
|
||||
|
||||
@@ -751,7 +815,8 @@ def create_shared_experts_from_config(
|
||||
in_dtype: torch.dtype,
|
||||
tp_size: int = 1,
|
||||
tp_rank: int = 0,
|
||||
device: torch.device | str | None = None,
|
||||
is_sequence_parallel: bool = False,
|
||||
device: torch.device | str | None = "cuda",
|
||||
) -> TestMLP | None:
|
||||
"""Create TestMLP for shared experts from config.
|
||||
|
||||
@@ -772,7 +837,7 @@ def create_shared_experts_from_config(
|
||||
s_w2 = shared_experts_config.w2
|
||||
|
||||
# Apply TP chunking if needed
|
||||
if tp_size > 1:
|
||||
if tp_size > 1 and not is_sequence_parallel:
|
||||
s_w1 = tp_chunk_gate_up(s_w1, tp_rank, tp_size, dim=1, device=device)
|
||||
s_w2 = chunk_by_rank(s_w2, tp_rank, tp_size, dim=0, device=device)
|
||||
else:
|
||||
@@ -920,6 +985,7 @@ def make_fused_moe_layer(
|
||||
routed_input_transform: torch.nn.Module | None = None,
|
||||
routed_output_transform: torch.nn.Module | None = None,
|
||||
pcp_size: int | None = 1,
|
||||
is_sequence_parallel: bool = False,
|
||||
) -> FusedMoE:
|
||||
quant_config, qw = make_quant_config(quantization, w1, w2, global_num_experts)
|
||||
|
||||
@@ -959,6 +1025,7 @@ def make_fused_moe_layer(
|
||||
enable_eplb=enable_eplb,
|
||||
num_redundant_experts=num_redundant_experts,
|
||||
has_bias=has_bias,
|
||||
is_sequence_parallel=is_sequence_parallel,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@@ -1004,7 +1071,6 @@ def make_fake_moe_layer(
|
||||
activation: str = "silu",
|
||||
indices_type: torch.dtype | None = None,
|
||||
expert_map: torch.Tensor | None = None,
|
||||
enable_eplb: bool = False,
|
||||
expert_load_view: torch.Tensor | None = None,
|
||||
logical_to_physical_map: torch.Tensor | None = None,
|
||||
logical_replica_count: torch.Tensor | None = None,
|
||||
@@ -1015,6 +1081,7 @@ def make_fake_moe_layer(
|
||||
tp_size: int = 1,
|
||||
dp_size: int = 1,
|
||||
ep_size: int = 1,
|
||||
is_sequence_parallel: bool = False,
|
||||
) -> Callable:
|
||||
quant_dtype = None
|
||||
activation = MoEActivation.from_str(activation)
|
||||
@@ -1022,7 +1089,6 @@ def make_fake_moe_layer(
|
||||
router = create_fused_moe_router(
|
||||
top_k=top_k,
|
||||
global_num_experts=global_num_experts,
|
||||
# eplb_state=None, # TODO
|
||||
renormalize=renormalize,
|
||||
use_grouped_topk=use_grouped_topk,
|
||||
num_expert_group=num_expert_group,
|
||||
@@ -1032,7 +1098,6 @@ def make_fake_moe_layer(
|
||||
routed_scaling_factor=routed_scaling_factor,
|
||||
e_score_correction_bias=e_score_correction_bias,
|
||||
num_fused_shared_experts=0, # TODO
|
||||
enable_eplb=enable_eplb,
|
||||
# TODO(bnell): once we can construct the MK at init time, we
|
||||
# can make this a value.
|
||||
indices_type_getter=lambda: indices_type,
|
||||
@@ -1046,7 +1111,8 @@ def make_fake_moe_layer(
|
||||
w2_s = None
|
||||
|
||||
shared_experts = create_shared_experts_from_config(
|
||||
shared_experts_config, in_dtype, 1, 0, "cuda"
|
||||
shared_experts_config,
|
||||
in_dtype,
|
||||
)
|
||||
|
||||
quant_config = FusedMoEQuantConfig.make(
|
||||
@@ -1117,7 +1183,7 @@ def make_fake_moe_layer(
|
||||
|
||||
|
||||
def _test_body_regular(
|
||||
moe_layer: Callable,
|
||||
moe_layer: FusedMoE,
|
||||
hidden_states: torch.Tensor,
|
||||
router_logits: torch.Tensor,
|
||||
vllm_config: VllmConfig,
|
||||
@@ -1134,7 +1200,7 @@ def _test_body_regular(
|
||||
num_tokens=num_tokens,
|
||||
num_tokens_across_dp=num_tokens_across_dp,
|
||||
):
|
||||
output = moe_layer(hidden_states, router_logits)
|
||||
output = sp_wrapper(moe_layer)(hidden_states, router_logits)
|
||||
|
||||
return baseline_output, output
|
||||
|
||||
@@ -1167,6 +1233,8 @@ def _test_body_eplb(
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
device = torch.accelerator.current_accelerator()
|
||||
|
||||
is_sequence_parallel = moe_layer.is_sequence_parallel
|
||||
|
||||
"""EPLB test body: compare output before and after expert weight rearrangement."""
|
||||
# Get "before" output with original weight arrangement
|
||||
with set_forward_context(
|
||||
@@ -1175,7 +1243,7 @@ def _test_body_eplb(
|
||||
num_tokens=num_tokens,
|
||||
num_tokens_across_dp=num_tokens_across_dp,
|
||||
):
|
||||
output_before = moe_layer(hidden_states, router_logits)
|
||||
output_before = sp_wrapper(moe_layer)(hidden_states, router_logits)
|
||||
|
||||
# Create a fresh FusedMoE layer with enable_eplb=True
|
||||
# Delete the original layer's registration so the constructor can
|
||||
@@ -1206,6 +1274,7 @@ def _test_body_eplb(
|
||||
gate=gate,
|
||||
routed_input_transform=routed_input_transform,
|
||||
routed_output_transform=routed_output_transform,
|
||||
is_sequence_parallel=is_sequence_parallel,
|
||||
)
|
||||
|
||||
if eplb_moe_layer._expert_map is not None:
|
||||
@@ -1267,7 +1336,7 @@ def _test_body_eplb(
|
||||
num_tokens=num_tokens,
|
||||
num_tokens_across_dp=num_tokens_across_dp,
|
||||
):
|
||||
output_after = eplb_moe_layer(hidden_states, router_logits)
|
||||
output_after = sp_wrapper(eplb_moe_layer)(hidden_states, router_logits)
|
||||
|
||||
return output_before, output_after
|
||||
|
||||
@@ -1275,11 +1344,12 @@ def _test_body_eplb(
|
||||
# TODO: make this take a MoETestConfig
|
||||
def _run_one_config(
|
||||
vllm_config: VllmConfig,
|
||||
ep_size: int,
|
||||
dp_size: int,
|
||||
tp_size: int,
|
||||
dp_rank: int,
|
||||
tp_rank: int,
|
||||
ep_size: int, # Expert parallel size (total across all ranks)
|
||||
dp_size: int, # Data parallel size (number of DP groups)
|
||||
tp_size: int, # Tensor parallel size OR sequence parallel size (when use_ep=True)
|
||||
dp_rank: int, # Current rank in data parallel dimension
|
||||
tp_rank: int, # Current rank in tensor/sequence parallel dimension
|
||||
is_sequence_parallel: bool, # Whether to use sequence parallelism
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
@@ -1293,15 +1363,22 @@ def _run_one_config(
|
||||
use_routed_input_transform: bool,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
set_random_seed(7)
|
||||
|
||||
"""Generic test loop that sets up environment and delegates to test_body_fn.
|
||||
|
||||
This function is called directly by test_moe_layer and test_moe_layer_eplb
|
||||
via parallel_launch_with_config, passing either _test_body_regular or
|
||||
_test_body_eplb as the test_body_fn parameter.
|
||||
Parameter Interpretation:
|
||||
- When is_sequence_parallel=False (standard TP or EP):
|
||||
* ep_size: Number of expert parallel ranks (or 1 if no EP)
|
||||
* tp_size: Number of tensor parallel ranks (or 1 if no TP)
|
||||
* Weights are chunked by ep_size (experts) and tp_size (tensors)
|
||||
|
||||
- When is_sequence_parallel=True (EP + sequence splitting):
|
||||
* ep_size: Number of expert parallel ranks (equals dp_size * tp_size)
|
||||
* tp_size: Number of ranks to split sequence across (becomes sp_size in FusedMoE)
|
||||
* Weights are chunked by ep_size (experts) but NOT by tp_size
|
||||
* Input sequences are chunked by tp_size (via sp_wrapper)
|
||||
"""
|
||||
world_size = tp_size * dp_size
|
||||
set_random_seed(7)
|
||||
|
||||
use_ep = ep_size > 1
|
||||
|
||||
assert vllm_config.parallel_config.enable_expert_parallel == use_ep
|
||||
@@ -1337,6 +1414,8 @@ def _run_one_config(
|
||||
routed_output_transform = test_data.routed_output_transform
|
||||
activation = "silu"
|
||||
|
||||
# Create baseline layer with FULL weights (no EP chunking)
|
||||
# Baseline represents the expected output using full model
|
||||
baseline_layer = make_fake_moe_layer(
|
||||
w1=w1,
|
||||
w2=w2,
|
||||
@@ -1354,26 +1433,45 @@ def _run_one_config(
|
||||
ep_size=ep_size,
|
||||
dp_size=dp_size,
|
||||
activation=activation,
|
||||
is_sequence_parallel=is_sequence_parallel,
|
||||
)
|
||||
|
||||
baseline_output = baseline_layer(hidden_states, router_logits)
|
||||
with set_current_vllm_config(vllm_config):
|
||||
# Compute baseline output with SP wrapper if needed
|
||||
# sp_wrapper handles sequence chunking/gathering for SP
|
||||
baseline_output = sp_wrapper(baseline_layer, is_sequence_parallel)(
|
||||
hidden_states, router_logits
|
||||
)
|
||||
|
||||
del baseline_layer
|
||||
torch.accelerator.empty_cache()
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
# Chunk weights for EP/TP (after baseline is created)
|
||||
# Chunk weights for EP BEFORE creating FusedMoE
|
||||
# FusedMoE uses EP-chunked weights and handles reductions internally
|
||||
if ep_size > 1:
|
||||
w1 = chunk_by_rank(w1, dp_rank, dp_size, dim=0, device=device)
|
||||
w2 = chunk_by_rank(w2, dp_rank, dp_size, dim=0, device=device)
|
||||
# Split experts across ranks (dimension 0 is the expert dimension)
|
||||
# When EP is enabled, use EP group rank and ep_size for chunking
|
||||
ep_rank = get_ep_group().rank_in_group
|
||||
w1 = chunk_by_rank(w1, ep_rank, ep_size, dim=0, device=device)
|
||||
w2 = chunk_by_rank(w2, ep_rank, ep_size, dim=0, device=device)
|
||||
|
||||
if tp_size > 1:
|
||||
# Chunk weights for TP (only if NOT doing sequence parallelism)
|
||||
# Sequence parallelism splits tokens/sequences, not weight tensors
|
||||
if tp_size > 1 and not is_sequence_parallel:
|
||||
w1 = tp_chunk_gate_up(w1, tp_rank, tp_size, dim=1, device=device)
|
||||
w2 = chunk_by_rank(w2, tp_rank, tp_size, dim=2, device=device)
|
||||
|
||||
# Setup shared experts if needed
|
||||
# In SP mode, shared experts should NOT be TP-chunked (same as routed experts)
|
||||
# tp_size is used for sequence splitting, not weight splitting
|
||||
shared_experts = create_shared_experts_from_config(
|
||||
shared_experts_config, in_dtype, tp_size, tp_rank, device
|
||||
shared_experts_config,
|
||||
in_dtype,
|
||||
tp_size,
|
||||
tp_rank,
|
||||
is_sequence_parallel,
|
||||
device,
|
||||
)
|
||||
|
||||
# Determine hidden size for MoE layer
|
||||
@@ -1399,14 +1497,17 @@ def _run_one_config(
|
||||
routed_input_transform=routed_input_transform,
|
||||
routed_output_transform=routed_output_transform,
|
||||
activation=activation,
|
||||
is_sequence_parallel=is_sequence_parallel,
|
||||
)
|
||||
|
||||
if moe_layer._expert_map is not None:
|
||||
moe_layer._expert_map = moe_layer._expert_map.to(device)
|
||||
|
||||
num_tokens = m
|
||||
# num_tokens_across_dp should have one entry per DP group, not per total rank
|
||||
# When EP is enabled, dp_size represents the number of DP groups
|
||||
num_tokens_across_dp = torch.tensor(
|
||||
[num_tokens] * world_size,
|
||||
[num_tokens] * dp_size,
|
||||
device=device,
|
||||
dtype=torch.int,
|
||||
)
|
||||
@@ -1448,7 +1549,7 @@ def _run_one_config(
|
||||
else:
|
||||
atol, rtol = 3.5e-2, 3.5e-2
|
||||
elif quantization in ("fp8", "fp8_blocked", "modelopt_fp8"):
|
||||
atol, rtol = 6e-2, 6e-2
|
||||
atol, rtol = 6.5e-2, 6.5e-2
|
||||
elif quantization == "modelopt_fp4":
|
||||
if k >= 2048:
|
||||
atol = rtol = 1e-1 + (k * 1e-4)
|
||||
@@ -1532,6 +1633,7 @@ def test_moe_layer_no_parallel(
|
||||
test_config.tp_size,
|
||||
0,
|
||||
0,
|
||||
False,
|
||||
test_config.m,
|
||||
test_config.n,
|
||||
test_config.k,
|
||||
@@ -1592,6 +1694,7 @@ def _parallel_worker(
|
||||
test_config.tp_size,
|
||||
dp_rank,
|
||||
tp_rank,
|
||||
test_config.is_sequence_parallel,
|
||||
test_config.m,
|
||||
test_config.n,
|
||||
test_config.k,
|
||||
@@ -1676,7 +1779,12 @@ def test_moe_layer(
|
||||
"""
|
||||
num_gpus = current_platform.device_count()
|
||||
world_size = tp_size * dp_size
|
||||
ep_size = 1 if not use_ep else world_size # or dp_size?
|
||||
# When use_ep=True: FusedMoEParallelConfig flattens tp_size across dp ranks
|
||||
# Result: ep_size = dp_size * pcp_size * tp_size
|
||||
# Since pcp_size=1 in these tests: ep_size = dp_size * tp_size = world_size
|
||||
# When use_ep=False: no expert parallelism, ep_size = 1
|
||||
ep_size = 1 if not use_ep else world_size
|
||||
|
||||
assert world_size > 1
|
||||
|
||||
# Check if enough GPUs available
|
||||
|
||||
@@ -10,7 +10,9 @@ import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.fused_moe import fused_topk
|
||||
from vllm.model_executor.layers.fused_moe.layer import determine_expert_map
|
||||
from vllm.model_executor.layers.fused_moe.expert_map_manager import (
|
||||
determine_expert_map,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
|
||||
moe_permute,
|
||||
moe_permute_unpermute_supported,
|
||||
|
||||
@@ -36,9 +36,11 @@ TOP_KS = [2, 4, 6]
|
||||
NUM_EXPERTS = [8, 16, 64]
|
||||
|
||||
|
||||
def setup_eplb_state(enable_eplb: bool, global_num_experts: int) -> EplbLayerState:
|
||||
def setup_eplb_state(
|
||||
enable_eplb: bool, global_num_experts: int
|
||||
) -> EplbLayerState | None:
|
||||
if not enable_eplb:
|
||||
return EplbLayerState()
|
||||
return None
|
||||
|
||||
# Initialize EPLB state with proper tensors for testing
|
||||
# For testing purposes, we use a simple 1:1 mapping (no redundant experts)
|
||||
@@ -349,7 +351,6 @@ def test_fused_topk(
|
||||
top_k=top_k,
|
||||
global_num_experts=global_num_experts,
|
||||
renormalize=renormalize,
|
||||
enable_eplb=enable_eplb,
|
||||
eplb_state=eplb_state,
|
||||
)
|
||||
|
||||
@@ -400,7 +401,6 @@ def test_fused_topk_bias(
|
||||
top_k=top_k,
|
||||
global_num_experts=global_num_experts,
|
||||
renormalize=renormalize,
|
||||
enable_eplb=enable_eplb,
|
||||
eplb_state=eplb_state,
|
||||
)
|
||||
|
||||
@@ -469,7 +469,6 @@ def test_grouped_topk(
|
||||
top_k=top_k,
|
||||
global_num_experts=global_num_experts,
|
||||
renormalize=renormalize,
|
||||
enable_eplb=enable_eplb,
|
||||
eplb_state=eplb_state,
|
||||
)
|
||||
|
||||
@@ -540,7 +539,6 @@ def test_custom(
|
||||
global_num_experts=global_num_experts,
|
||||
custom_routing_function=custom_routing_function,
|
||||
renormalize=renormalize,
|
||||
enable_eplb=enable_eplb,
|
||||
eplb_state=eplb_state,
|
||||
)
|
||||
|
||||
@@ -580,7 +578,6 @@ def test_custom(
|
||||
# router = create_fused_moe_router(
|
||||
# top_k=top_k,
|
||||
# global_num_experts=global_num_experts,
|
||||
# enable_eplb=enable_eplb,
|
||||
# eplb_state=eplb_state,
|
||||
# )
|
||||
|
||||
|
||||
@@ -183,6 +183,7 @@ def test_zero_expert_moe_output_decomposition(zero_expert_moe, num_tokens):
|
||||
x=hidden_states,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
shared_experts=None,
|
||||
shared_experts_input=None,
|
||||
)
|
||||
|
||||
|
||||
@@ -17,13 +17,13 @@ from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantConfig,
|
||||
RoutingMethodType,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
|
||||
TritonExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
|
||||
from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
|
||||
BatchedTritonExperts,
|
||||
NaiveBatchedExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.triton_moe import (
|
||||
TritonExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import (
|
||||
fused_experts,
|
||||
)
|
||||
|
||||
@@ -143,6 +143,15 @@ def test_models(
|
||||
# in parts of the operators
|
||||
pytest.skip(f"Skipping '{model}' model test with AITER kernel.")
|
||||
|
||||
if model == "bigcode/starcoder2-3b":
|
||||
# Replace example.txt's Test1 (an NL prompt) with a code prompt:
|
||||
# starcoder2-3b is a code model, so NL prompts give near-uniform
|
||||
# digit logits where HF<->vLLM bf16 drift can reorder top-K.
|
||||
example_prompts = list(example_prompts)
|
||||
example_prompts[1] = (
|
||||
"def add(a, b):\n return a + b\n\ndef sub(a, b):\n return a - "
|
||||
)
|
||||
|
||||
with hf_runner(model) as hf_model:
|
||||
hf_outputs = hf_model.generate_greedy_logprobs_limit(
|
||||
example_prompts, max_tokens, num_logprobs
|
||||
|
||||
@@ -42,6 +42,18 @@ def qwen_vl_chat_template(content: str) -> str:
|
||||
|
||||
|
||||
MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
|
||||
"qwen2_5_vl": VitCudagraphTestConfig(
|
||||
model="Qwen/Qwen2.5-VL-3B-Instruct",
|
||||
image_prompt=qwen_vl_chat_template(
|
||||
"<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
|
||||
),
|
||||
video_prompt=qwen_vl_chat_template(
|
||||
"<|vision_start|><|video_pad|><|vision_end|>"
|
||||
"Describe this video in one sentence."
|
||||
),
|
||||
needs_video_metadata=False,
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
"qwen3_vl": VitCudagraphTestConfig(
|
||||
model="Qwen/Qwen3-VL-2B-Instruct",
|
||||
image_prompt=qwen_vl_chat_template(
|
||||
@@ -54,8 +66,8 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
|
||||
needs_video_metadata=True,
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
"qwen2_5_vl": VitCudagraphTestConfig(
|
||||
model="Qwen/Qwen2.5-VL-3B-Instruct",
|
||||
"qwen3_5": VitCudagraphTestConfig(
|
||||
model="Qwen/Qwen3.5-0.8B",
|
||||
image_prompt=qwen_vl_chat_template(
|
||||
"<|vision_start|><|image_pad|><|vision_end|>What is in this image?"
|
||||
),
|
||||
@@ -63,7 +75,7 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
|
||||
"<|vision_start|><|video_pad|><|vision_end|>"
|
||||
"Describe this video in one sentence."
|
||||
),
|
||||
needs_video_metadata=False,
|
||||
needs_video_metadata=True,
|
||||
marks=[pytest.mark.core_model],
|
||||
),
|
||||
}
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for Step3-VL precomputed image embedding inputs."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.models.step3_vl import (
|
||||
Step3VLForConditionalGeneration,
|
||||
Step3VLImageEmbeddingInputs,
|
||||
)
|
||||
|
||||
|
||||
class _FakeStep3VL:
|
||||
@staticmethod
|
||||
def _process_image_features(image_features: torch.Tensor) -> torch.Tensor:
|
||||
return image_features
|
||||
|
||||
|
||||
def test_image_embedding_inputs_construction():
|
||||
"""Step3VLImageEmbeddingInputs should store embeddings in the data field."""
|
||||
image_embeds = torch.randn(2, 16, 64)
|
||||
|
||||
inputs = Step3VLImageEmbeddingInputs(
|
||||
type="image_embeds",
|
||||
data=image_embeds,
|
||||
)
|
||||
|
||||
assert inputs["type"] == "image_embeds"
|
||||
assert torch.equal(inputs["data"], image_embeds)
|
||||
assert torch.equal(inputs.data, image_embeds)
|
||||
|
||||
|
||||
def test_image_embedding_inputs_validation_rejects_wrong_rank():
|
||||
"""Validation should reject tensors with wrong rank."""
|
||||
with pytest.raises(ValueError, match="rank"):
|
||||
Step3VLImageEmbeddingInputs(
|
||||
type="image_embeds",
|
||||
data=torch.randn(16, 64),
|
||||
)
|
||||
|
||||
|
||||
def test_process_image_embeds_does_not_require_pixel_input_fields():
|
||||
"""The image_embeds branch should not reference patch pixel metadata."""
|
||||
image_embeds = torch.randn(2, 4, 8)
|
||||
image_input = Step3VLImageEmbeddingInputs(
|
||||
type="image_embeds",
|
||||
data=image_embeds,
|
||||
)
|
||||
|
||||
outputs = Step3VLForConditionalGeneration._process_image_input(
|
||||
_FakeStep3VL(),
|
||||
image_input,
|
||||
)
|
||||
|
||||
assert len(outputs) == 2
|
||||
assert torch.equal(outputs[0], image_embeds[0])
|
||||
assert torch.equal(outputs[1], image_embeds[1])
|
||||
@@ -1098,6 +1098,10 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
},
|
||||
trust_remote_code=True,
|
||||
),
|
||||
"MiniCPMV4_6ForConditionalGeneration": _HfExamplesInfo(
|
||||
"openbmb/MiniCPM-V-4_6",
|
||||
min_transformers_version="5.7.0",
|
||||
),
|
||||
"MiniMaxVL01ForConditionalGeneration": _HfExamplesInfo(
|
||||
"MiniMaxAI/MiniMax-VL-01",
|
||||
trust_remote_code=True,
|
||||
@@ -1489,6 +1493,21 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
|
||||
"Qwen/Qwen3-VL-8B-Instruct",
|
||||
speculative_model="taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
|
||||
),
|
||||
# [PEagle]
|
||||
"PEagleDraftModel": _HfExamplesInfo(
|
||||
"Qwen/Qwen3-8B",
|
||||
trust_remote_code=True,
|
||||
speculative_model="nm-testing/qwen3-8b-peagle-speculators",
|
||||
tokenizer="Qwen/Qwen3-8B",
|
||||
use_original_num_layers=True,
|
||||
),
|
||||
"PeagleLlamaForCausalLM": _HfExamplesInfo(
|
||||
"Qwen/Qwen3-8B",
|
||||
trust_remote_code=True,
|
||||
speculative_model="nm-testing/qwen3-8b-peagle-speculators",
|
||||
tokenizer="Qwen/Qwen3-8B",
|
||||
use_original_num_layers=True,
|
||||
),
|
||||
# [MTP]
|
||||
"DeepSeekMTPModel": _HfExamplesInfo(
|
||||
"luccafong/deepseek_mtp_main_random",
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# test_audio.py
|
||||
import math
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
@@ -45,7 +46,6 @@ def test_resample_audio_scipy(dummy_audio):
|
||||
assert np.all(out_same == dummy_audio)
|
||||
|
||||
|
||||
@pytest.mark.xfail(reason="resample_audio_scipy is buggy for non-integer ratios")
|
||||
def test_resample_audio_scipy_non_integer_ratio(dummy_audio):
|
||||
out = resample_audio_scipy(dummy_audio, orig_sr=5, target_sr=3)
|
||||
|
||||
@@ -56,6 +56,26 @@ def test_resample_audio_scipy_non_integer_ratio(dummy_audio):
|
||||
assert np.isfinite(out).all()
|
||||
|
||||
|
||||
def test_resample_audio_scipy_non_divisible_sample_rates():
|
||||
audio = np.arange(441, dtype=float)
|
||||
out = resample_audio_scipy(audio, orig_sr=44100, target_sr=16000)
|
||||
|
||||
expected_len = math.ceil(len(audio) * 16000 / 44100)
|
||||
assert len(out) == expected_len
|
||||
|
||||
assert isinstance(out, np.ndarray)
|
||||
assert np.isfinite(out).all()
|
||||
|
||||
|
||||
def test_resample_audio_scipy_resamples_last_axis_for_multichannel():
|
||||
audio = np.arange(2 * 441, dtype=float).reshape(2, 441)
|
||||
out = resample_audio_scipy(audio, orig_sr=44100, target_sr=16000)
|
||||
|
||||
expected_len = math.ceil(audio.shape[-1] * 16000 / 44100)
|
||||
assert out.shape == (2, expected_len)
|
||||
assert np.isfinite(out).all()
|
||||
|
||||
|
||||
def test_audio_resampler_pyav_calls_resample(dummy_audio):
|
||||
resampler = AudioResampler(target_sr=22050, method="pyav")
|
||||
with patch("vllm.multimodal.audio.resample_audio_pyav") as mock_resample:
|
||||
|
||||
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
|
||||
CompressedTensorsConfig,
|
||||
CompressedTensorsLinearMethod,
|
||||
CompressedTensorsW4A4Fp4,
|
||||
CompressedTensorsW4A4Mxfp4,
|
||||
CompressedTensorsW4A8Fp8,
|
||||
CompressedTensorsW4A16Fp4,
|
||||
CompressedTensorsW8A8Fp8,
|
||||
@@ -689,3 +690,31 @@ def test_compressed_tensors_mxfp8_moe_setup(vllm_runner):
|
||||
llm.apply_model(check_model)
|
||||
output = llm.generate_greedy("Hello my name is", max_tokens=4)
|
||||
assert output
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda() or not current_platform.has_device_capability(80),
|
||||
reason="MXFP4 requires ampere or newer",
|
||||
)
|
||||
def test_compressed_tensors_mxfp4(vllm_runner):
|
||||
model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-MXFP4"
|
||||
with vllm_runner(model_path, enforce_eager=True) as llm:
|
||||
|
||||
def check_model(model):
|
||||
layer = model.model.layers[0]
|
||||
|
||||
qkv_proj = layer.self_attn.qkv_proj
|
||||
o_proj = layer.self_attn.o_proj
|
||||
gate_up_proj = layer.mlp.gate_up_proj
|
||||
down_proj = layer.mlp.down_proj
|
||||
|
||||
for proj in (qkv_proj, o_proj, gate_up_proj, down_proj):
|
||||
assert isinstance(proj.quant_method, CompressedTensorsLinearMethod)
|
||||
assert isinstance(proj.scheme, CompressedTensorsW4A4Mxfp4)
|
||||
|
||||
# Verify group size
|
||||
assert proj.scheme.group_size == 32
|
||||
|
||||
llm.apply_model(check_model)
|
||||
output = llm.generate_greedy("Hello my name is", max_tokens=4)
|
||||
assert output
|
||||
|
||||
@@ -14,7 +14,7 @@ import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.has_device_capability(90):
|
||||
if not current_platform.has_device_capability(90) or current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"Machete W4A16 requires Hopper (sm_90).",
|
||||
allow_module_level=True,
|
||||
|
||||
@@ -43,11 +43,16 @@ from vllm.v1.kv_cache_interface import (
|
||||
KVCacheConfig,
|
||||
KVCacheGroupSpec,
|
||||
KVCacheSpec,
|
||||
KVCacheSpecKind,
|
||||
KVCacheTensor,
|
||||
MambaSpec,
|
||||
MLAAttentionSpec,
|
||||
SinkFullAttentionSpec,
|
||||
SlidingWindowMLASpec,
|
||||
SlidingWindowSpec,
|
||||
UniformTypeKVCacheSpecs,
|
||||
get_kv_cache_spec_kind,
|
||||
get_kv_cache_spec_sliding_window,
|
||||
)
|
||||
from vllm.v1.metrics.stats import CachingMetrics, PrefixCacheStats
|
||||
from vllm.v1.request import Request
|
||||
@@ -1865,6 +1870,149 @@ def new_mla_spec(cache_dtype_str=None):
|
||||
)
|
||||
|
||||
|
||||
def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses():
|
||||
assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION
|
||||
|
||||
sliding_window_mla_spec = SlidingWindowMLASpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=576,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128,
|
||||
)
|
||||
assert (
|
||||
get_kv_cache_spec_kind(sliding_window_mla_spec)
|
||||
== KVCacheSpecKind.SLIDING_WINDOW_MLA
|
||||
)
|
||||
|
||||
sink_full_attention_spec = SinkFullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=64,
|
||||
dtype=torch.float32,
|
||||
sink_len=4,
|
||||
)
|
||||
assert (
|
||||
get_kv_cache_spec_kind(sink_full_attention_spec)
|
||||
== KVCacheSpecKind.SINK_FULL_ATTENTION
|
||||
)
|
||||
|
||||
|
||||
def test_get_kv_cache_spec_kind_unwraps_uniform_type_specs():
|
||||
uniform_mla_spec = UniformTypeKVCacheSpecs(
|
||||
block_size=16,
|
||||
kv_cache_specs={
|
||||
"layer_1": new_mla_spec(),
|
||||
"layer_2": new_mla_spec(cache_dtype_str="fp8"),
|
||||
},
|
||||
)
|
||||
assert get_kv_cache_spec_kind(uniform_mla_spec) == KVCacheSpecKind.MLA_ATTENTION
|
||||
|
||||
uniform_swa_mla_spec = UniformTypeKVCacheSpecs(
|
||||
block_size=16,
|
||||
kv_cache_specs={
|
||||
"layer_1": SlidingWindowMLASpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=576,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128,
|
||||
),
|
||||
"layer_2": SlidingWindowMLASpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=1024,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128,
|
||||
),
|
||||
},
|
||||
)
|
||||
assert (
|
||||
get_kv_cache_spec_kind(uniform_swa_mla_spec)
|
||||
== KVCacheSpecKind.SLIDING_WINDOW_MLA
|
||||
)
|
||||
|
||||
|
||||
def test_get_kv_cache_spec_kind_unknown_for_mixed_uniform_type_specs():
|
||||
uniform_mixed_spec = UniformTypeKVCacheSpecs(
|
||||
block_size=16,
|
||||
kv_cache_specs={
|
||||
"layer_1": new_mla_spec(),
|
||||
"layer_2": SlidingWindowMLASpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=576,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128,
|
||||
),
|
||||
},
|
||||
)
|
||||
assert get_kv_cache_spec_kind(uniform_mixed_spec) == KVCacheSpecKind.UNKNOWN
|
||||
|
||||
|
||||
def test_get_kv_cache_spec_sliding_window_reads_windowed_specs():
|
||||
full_attention_spec = FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=64,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
sliding_window_spec = SlidingWindowSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=64,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128,
|
||||
)
|
||||
|
||||
assert get_kv_cache_spec_sliding_window(full_attention_spec) is None
|
||||
assert get_kv_cache_spec_sliding_window(sliding_window_spec) == 128
|
||||
|
||||
|
||||
def test_get_kv_cache_spec_sliding_window_unwraps_uniform_type_specs():
|
||||
uniform_window_spec = UniformTypeKVCacheSpecs(
|
||||
block_size=16,
|
||||
kv_cache_specs={
|
||||
"layer_1": SlidingWindowSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=64,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128,
|
||||
),
|
||||
"layer_2": SlidingWindowSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=2,
|
||||
head_size=64,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128,
|
||||
),
|
||||
},
|
||||
)
|
||||
mixed_window_spec = UniformTypeKVCacheSpecs(
|
||||
block_size=16,
|
||||
kv_cache_specs={
|
||||
"layer_1": SlidingWindowSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=64,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128,
|
||||
),
|
||||
"layer_2": SlidingWindowSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=1,
|
||||
head_size=64,
|
||||
dtype=torch.float32,
|
||||
sliding_window=256,
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
assert get_kv_cache_spec_sliding_window(uniform_window_spec) == 128
|
||||
assert get_kv_cache_spec_sliding_window(mixed_window_spec) is None
|
||||
|
||||
|
||||
def test_merge_mla_spec():
|
||||
kv_cache_specs = [
|
||||
new_mla_spec(),
|
||||
|
||||
@@ -8,6 +8,7 @@ from collections.abc import Callable
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm.v1.core.kv_cache_manager as kv_cache_manager
|
||||
import vllm.v1.core.kv_cache_utils as kv_cache_utils
|
||||
from vllm.distributed.kv_events import AllBlocksCleared, BlockRemoved, BlockStored
|
||||
from vllm.lora.request import LoRARequest
|
||||
@@ -35,6 +36,7 @@ from vllm.v1.kv_cache_interface import (
|
||||
FullAttentionSpec,
|
||||
KVCacheConfig,
|
||||
KVCacheGroupSpec,
|
||||
KVCacheSpecKind,
|
||||
MambaSpec,
|
||||
SlidingWindowSpec,
|
||||
)
|
||||
@@ -1933,6 +1935,7 @@ def test_kv_cache_events(blocks_to_cache: int):
|
||||
== len(manager.block_pool.cached_block_hash_to_block)
|
||||
)
|
||||
assert len(block.token_ids) == block.block_size * len(block.block_hashes)
|
||||
assert block.kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
|
||||
assert len(manager.block_pool.kv_event_queue) == 0
|
||||
|
||||
stored_block_hash = block.block_hashes
|
||||
@@ -1946,6 +1949,7 @@ def test_kv_cache_events(blocks_to_cache: int):
|
||||
events = manager.take_events()
|
||||
|
||||
for blocks in events[:-1]:
|
||||
assert isinstance(blocks, BlockRemoved)
|
||||
assert blocks.block_hashes[0] in stored_block_hash
|
||||
assert len(events) == blocks_to_cache + 1
|
||||
assert isinstance(events[-2], BlockRemoved)
|
||||
@@ -2022,6 +2026,8 @@ def test_null_parent_block_hash():
|
||||
]
|
||||
assert event.block_hashes == expected_new_hashes
|
||||
assert event.group_idx == kv_cache_group_id
|
||||
assert event.kv_cache_spec_kind is None
|
||||
assert event.kv_cache_spec_sliding_window is None
|
||||
|
||||
# Ensure we didn't accidentally assign a hash to the null block.
|
||||
assert pool.null_block.block_hash is None
|
||||
@@ -2095,12 +2101,14 @@ def test_block_stored_event_group_idx(group_id: int):
|
||||
block_size = 4
|
||||
num_tokens = block_size * 2
|
||||
|
||||
pool = BlockPool(
|
||||
num_gpu_blocks=5,
|
||||
manager = KVCacheManager(
|
||||
make_kv_cache_config_three_types(block_size, num_blocks=5),
|
||||
max_model_len=8192,
|
||||
enable_caching=True,
|
||||
hash_block_size=block_size,
|
||||
enable_kv_cache_events=True,
|
||||
hash_block_size=block_size,
|
||||
)
|
||||
pool = manager.block_pool
|
||||
|
||||
req = make_request(
|
||||
"req_grp_idx",
|
||||
@@ -2119,10 +2127,26 @@ def test_block_stored_event_group_idx(group_id: int):
|
||||
kv_cache_group_id=group_id,
|
||||
)
|
||||
|
||||
events = pool.take_events()
|
||||
events = manager.take_events()
|
||||
assert len(events) == 1
|
||||
assert isinstance(events[0], BlockStored)
|
||||
assert events[0].group_idx == group_id
|
||||
assert (
|
||||
events[0].kv_cache_spec_kind
|
||||
== [
|
||||
KVCacheSpecKind.FULL_ATTENTION.value,
|
||||
KVCacheSpecKind.SLIDING_WINDOW.value,
|
||||
KVCacheSpecKind.MAMBA.value,
|
||||
][group_id]
|
||||
)
|
||||
assert (
|
||||
events[0].kv_cache_spec_sliding_window
|
||||
== [
|
||||
None,
|
||||
2 * block_size,
|
||||
None,
|
||||
][group_id]
|
||||
)
|
||||
|
||||
|
||||
def test_block_stored_event_group_idx_multiple_groups():
|
||||
@@ -2137,13 +2161,38 @@ def test_block_stored_event_group_idx_multiple_groups():
|
||||
block_size = 4
|
||||
num_tokens = block_size * 2
|
||||
|
||||
# null block + 4 usable (2 per group)
|
||||
pool = BlockPool(
|
||||
num_gpu_blocks=5,
|
||||
manager = KVCacheManager(
|
||||
KVCacheConfig(
|
||||
num_blocks=5,
|
||||
kv_cache_tensors=[],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(
|
||||
["layer1"],
|
||||
FullAttentionSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
),
|
||||
),
|
||||
KVCacheGroupSpec(
|
||||
["layer2"],
|
||||
SlidingWindowSpec(
|
||||
block_size=block_size,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
sliding_window=128,
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
max_model_len=8192,
|
||||
enable_caching=True,
|
||||
hash_block_size=block_size,
|
||||
enable_kv_cache_events=True,
|
||||
hash_block_size=block_size,
|
||||
)
|
||||
pool = manager.block_pool
|
||||
|
||||
req = make_request(
|
||||
"req_multi_grp",
|
||||
@@ -2174,12 +2223,52 @@ def test_block_stored_event_group_idx_multiple_groups():
|
||||
kv_cache_group_id=1,
|
||||
)
|
||||
|
||||
events = pool.take_events()
|
||||
events = manager.take_events()
|
||||
assert len(events) == 2
|
||||
assert isinstance(events[0], BlockStored)
|
||||
assert events[0].group_idx == 0
|
||||
assert events[0].kv_cache_spec_kind == KVCacheSpecKind.FULL_ATTENTION.value
|
||||
assert events[0].kv_cache_spec_sliding_window is None
|
||||
assert isinstance(events[1], BlockStored)
|
||||
assert events[1].group_idx == 1
|
||||
assert events[1].kv_cache_spec_kind == KVCacheSpecKind.SLIDING_WINDOW.value
|
||||
assert events[1].kv_cache_spec_sliding_window == 128
|
||||
|
||||
|
||||
def test_block_stored_event_group_idx_out_of_bounds(monkeypatch):
|
||||
"""Out-of-range group_idx events are returned without metadata annotation."""
|
||||
block_size = 4
|
||||
manager = KVCacheManager(
|
||||
make_kv_cache_config(block_size, num_blocks=5),
|
||||
max_model_len=8192,
|
||||
enable_caching=True,
|
||||
enable_kv_cache_events=True,
|
||||
hash_block_size=block_size,
|
||||
)
|
||||
event = BlockStored(
|
||||
block_hashes=[1],
|
||||
parent_block_hash=None,
|
||||
token_ids=list(range(block_size)),
|
||||
block_size=block_size,
|
||||
lora_id=None,
|
||||
medium=None,
|
||||
lora_name=None,
|
||||
group_idx=1,
|
||||
)
|
||||
manager.block_pool.kv_event_queue.append(event)
|
||||
warnings = []
|
||||
|
||||
def collect_warning(message, *args, **kwargs):
|
||||
del kwargs
|
||||
warnings.append(message % args if args else message)
|
||||
|
||||
monkeypatch.setattr(kv_cache_manager.logger, "warning", collect_warning)
|
||||
events = manager.take_events()
|
||||
|
||||
assert events == [event]
|
||||
assert event.kv_cache_spec_kind is None
|
||||
assert event.kv_cache_spec_sliding_window is None
|
||||
assert warnings == ["Group index `1` not in KV cache metadata"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_id", [0, 1, 2])
|
||||
|
||||
@@ -429,7 +429,7 @@ def _logprobs_match(
|
||||
and lps_a.keys() == lps_b.keys()
|
||||
and all(
|
||||
a.decoded_token == b.decoded_token
|
||||
and a.rank == b.rank
|
||||
and a.rank == pytest.approx(b.rank, rel=0.005)
|
||||
and a.logprob == pytest.approx(b.logprob, rel=rel_tol, abs=abs_tol)
|
||||
for a, b in ((lps_a[x], lps_b[x]) for x in lps_a)
|
||||
)
|
||||
|
||||
@@ -79,10 +79,11 @@ def _get_test_sampling_params(
|
||||
structured_outputs: bool = False,
|
||||
) -> tuple[list[SamplingParams], list[int]]:
|
||||
"""Generate random sampling params for a batch."""
|
||||
rng = random.Random(seed)
|
||||
|
||||
def get_mostly_n_gt1() -> int:
|
||||
r"""Mostly n \in [2,20], ~1/3 n=1"""
|
||||
x = random.randint(0, 28)
|
||||
x = rng.randint(0, 28)
|
||||
if x < 10:
|
||||
return 1
|
||||
else:
|
||||
|
||||
@@ -627,7 +627,7 @@ def _make_scheduler_with_lookup(
|
||||
return scheduler
|
||||
|
||||
|
||||
_EMPTY_REQ_CTX = ReqContext()
|
||||
_EMPTY_REQ_CTX = ReqContext(req_id="")
|
||||
|
||||
|
||||
class TestMaximalPrefixLookup:
|
||||
|
||||
@@ -459,9 +459,6 @@ class RequestRunner:
|
||||
self.worker_connector.bind_connector_metadata(kv_connector_metadata)
|
||||
self.worker_connector.start_load_kv(self._dummy_ctx)
|
||||
|
||||
if scheduler_output.total_num_scheduled_tokens > 0:
|
||||
self.worker_connector.wait_for_save()
|
||||
|
||||
if complete_transfers:
|
||||
self.offloading_spec.complete_transfers()
|
||||
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( # noqa: E501
|
||||
ExampleConnectorMetadata,
|
||||
)
|
||||
@@ -38,11 +40,20 @@ def test_kv_connector_mixin_clears_metadata():
|
||||
vllm_config.kv_transfer_config.kv_role = "kv_both"
|
||||
vllm_config.kv_transfer_config.kv_connector_extra_config["name"] = "unit"
|
||||
|
||||
# Initialize the global connector instance
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
|
||||
)
|
||||
ensure_kv_transfer_initialized(vllm_config, kv_cache_config)
|
||||
# Initialize the global connector instance.
|
||||
# kv_transfer init now syncs engine_id across TP, so unit tests need
|
||||
# a minimal mocked TP group.
|
||||
mock_tp_group = MagicMock()
|
||||
mock_tp_group.broadcast_object.side_effect = lambda value, src=0: value
|
||||
|
||||
with patch(
|
||||
"vllm.distributed.parallel_state.get_tp_group",
|
||||
return_value=mock_tp_group,
|
||||
):
|
||||
ensure_kv_transfer_initialized(vllm_config, kv_cache_config)
|
||||
|
||||
try:
|
||||
# Minimal scheduler output with empty metadata; mixin should still
|
||||
|
||||
@@ -0,0 +1,258 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from vllm.config import set_current_vllm_config
|
||||
from vllm.distributed.kv_events import BlockStored
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
|
||||
KVConnectorRole,
|
||||
)
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
|
||||
connector,
|
||||
worker,
|
||||
)
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
|
||||
MooncakeStoreConnectorMetadata,
|
||||
)
|
||||
from vllm.v1.outputs import KVConnectorOutput
|
||||
|
||||
from .utils import create_vllm_config
|
||||
|
||||
|
||||
def _make_vllm_config():
|
||||
return create_vllm_config(
|
||||
kv_connector="MooncakeStoreConnector",
|
||||
kv_role="kv_both",
|
||||
)
|
||||
|
||||
|
||||
def _make_block_stored() -> BlockStored:
|
||||
return BlockStored(
|
||||
block_hashes=[b"hash"],
|
||||
parent_block_hash=None,
|
||||
token_ids=[1, 2, 3],
|
||||
block_size=16,
|
||||
lora_id=None,
|
||||
medium="cpu",
|
||||
lora_name=None,
|
||||
)
|
||||
|
||||
|
||||
def test_scheduler_role_initializes_store_scheduler_only():
|
||||
vllm_config = _make_vllm_config()
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreScheduler"
|
||||
) as mock_scheduler,
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreWorker"
|
||||
) as mock_worker,
|
||||
):
|
||||
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
|
||||
|
||||
mock_scheduler.assert_called_once_with(vllm_config)
|
||||
mock_worker.assert_not_called()
|
||||
assert conn.connector_scheduler is mock_scheduler.return_value
|
||||
assert conn.connector_worker is None
|
||||
|
||||
|
||||
def test_worker_role_initializes_store_worker_on_rank0():
|
||||
vllm_config = _make_vllm_config()
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreScheduler"
|
||||
) as mock_scheduler,
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreWorker"
|
||||
) as mock_worker,
|
||||
):
|
||||
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
|
||||
|
||||
mock_scheduler.assert_not_called()
|
||||
mock_worker.assert_called_once_with(vllm_config)
|
||||
assert conn.connector_scheduler is None
|
||||
assert conn.connector_worker is mock_worker.return_value
|
||||
|
||||
|
||||
def test_worker_role_initializes_on_nonzero_rank():
|
||||
vllm_config = _make_vllm_config()
|
||||
vllm_config.parallel_config.rank = 1
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreWorker"
|
||||
) as mock_worker,
|
||||
):
|
||||
connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
|
||||
|
||||
mock_worker.assert_called_once_with(vllm_config)
|
||||
|
||||
|
||||
def test_lookup_rpc_path_uses_data_parallel_index_in_dense_dp():
|
||||
vllm_config = _make_vllm_config()
|
||||
vllm_config.parallel_config.data_parallel_rank = 0
|
||||
vllm_config.parallel_config.data_parallel_index = 3
|
||||
|
||||
path = worker.get_zmq_rpc_path_lookup(vllm_config)
|
||||
|
||||
assert path.endswith("_dp_rank3")
|
||||
|
||||
|
||||
def test_lookup_rpc_path_uses_local_rank_when_local_engines_only():
|
||||
vllm_config = _make_vllm_config()
|
||||
vllm_config.parallel_config.data_parallel_index = 7
|
||||
vllm_config.parallel_config.data_parallel_rank_local = 1
|
||||
vllm_config.parallel_config.data_parallel_hybrid_lb = True
|
||||
|
||||
path = worker.get_zmq_rpc_path_lookup(vllm_config)
|
||||
|
||||
assert path.endswith("_dp_rank1")
|
||||
|
||||
|
||||
def test_worker_methods_delegate_to_store_worker():
|
||||
vllm_config = _make_vllm_config()
|
||||
kv_caches = {"layer0": MagicMock()}
|
||||
metadata = MooncakeStoreConnectorMetadata(set(), set())
|
||||
finished_req_ids = {"req-1"}
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreWorker"
|
||||
) as mock_worker_cls,
|
||||
):
|
||||
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
|
||||
|
||||
worker_inst = mock_worker_cls.return_value
|
||||
worker_inst.get_finished.return_value = ({"req-1"}, {"req-2"})
|
||||
conn.bind_connector_metadata(metadata)
|
||||
|
||||
conn.register_kv_caches(kv_caches)
|
||||
result = conn.get_finished(finished_req_ids)
|
||||
|
||||
worker_inst.register_kv_caches.assert_called_once_with(kv_caches)
|
||||
worker_inst.get_finished.assert_called_once_with(finished_req_ids, metadata)
|
||||
assert result == ({"req-1"}, {"req-2"})
|
||||
|
||||
|
||||
def test_get_kv_connector_kv_cache_events_returns_none_when_empty():
|
||||
vllm_config = _make_vllm_config()
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreWorker"
|
||||
) as mock_worker_cls,
|
||||
):
|
||||
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
|
||||
|
||||
mock_worker_cls.return_value.get_kv_events.return_value = []
|
||||
assert conn.get_kv_connector_kv_cache_events() is None
|
||||
|
||||
|
||||
def test_get_kv_connector_kv_cache_events_wraps_worker_events():
|
||||
vllm_config = _make_vllm_config()
|
||||
event = _make_block_stored()
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreWorker"
|
||||
) as mock_worker_cls,
|
||||
):
|
||||
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
|
||||
|
||||
mock_worker_cls.return_value.get_kv_events.return_value = [event]
|
||||
kv_events = conn.get_kv_connector_kv_cache_events()
|
||||
|
||||
assert isinstance(kv_events, connector.MooncakeStoreKVEvents)
|
||||
assert kv_events.get_number_of_workers() == 1
|
||||
assert kv_events.get_all_events() == [event]
|
||||
|
||||
|
||||
def test_prefer_cross_layer_blocks_from_config():
|
||||
# Default: disabled
|
||||
vllm_config = _make_vllm_config()
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreScheduler"
|
||||
),
|
||||
):
|
||||
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
|
||||
assert conn.prefer_cross_layer_blocks is False
|
||||
|
||||
# Enabled via config
|
||||
vllm_config_enabled = create_vllm_config(
|
||||
kv_connector="MooncakeStoreConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config={"enable_cross_layers_blocks": "true"},
|
||||
)
|
||||
with (
|
||||
set_current_vllm_config(vllm_config_enabled),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreScheduler"
|
||||
),
|
||||
):
|
||||
conn_enabled = connector.MooncakeStoreConnector(
|
||||
vllm_config_enabled, KVConnectorRole.SCHEDULER
|
||||
)
|
||||
assert conn_enabled.prefer_cross_layer_blocks is True
|
||||
|
||||
|
||||
def test_register_cross_layers_kv_cache_delegates_to_worker():
|
||||
vllm_config = _make_vllm_config()
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreWorker"
|
||||
) as mock_worker_cls,
|
||||
):
|
||||
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.WORKER)
|
||||
|
||||
fake_tensor = MagicMock()
|
||||
fake_backend = MagicMock()
|
||||
conn.register_cross_layers_kv_cache(fake_tensor, fake_backend)
|
||||
|
||||
worker_inst = mock_worker_cls.return_value
|
||||
worker_inst.register_cross_layers_kv_caches.assert_called_once_with(fake_tensor)
|
||||
|
||||
|
||||
def test_update_connector_output_and_take_events():
|
||||
vllm_config = _make_vllm_config()
|
||||
event = _make_block_stored()
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"connector.MooncakeStoreScheduler"
|
||||
),
|
||||
):
|
||||
conn = connector.MooncakeStoreConnector(vllm_config, KVConnectorRole.SCHEDULER)
|
||||
|
||||
kv_events = connector.MooncakeStoreKVEvents(num_workers=1)
|
||||
kv_events.add_events([event])
|
||||
conn.update_connector_output(KVConnectorOutput(kv_cache_events=kv_events))
|
||||
|
||||
assert conn._kv_cache_events is kv_events
|
||||
assert list(conn.take_events()) == [event]
|
||||
assert conn._kv_cache_events is None
|
||||
@@ -0,0 +1,300 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import threading
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
|
||||
worker,
|
||||
)
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import ( # noqa: E501
|
||||
ChunkedTokenDatabase,
|
||||
KeyMetadata,
|
||||
ReqMeta,
|
||||
)
|
||||
|
||||
|
||||
def _make_store_sending_thread(
|
||||
store: MagicMock,
|
||||
) -> worker.KVCacheStoreSendingThread:
|
||||
token_database = ChunkedTokenDatabase(
|
||||
KeyMetadata("test-model", 0, 0, 0, 0), block_size=16
|
||||
)
|
||||
token_database.set_kv_caches_base_addr([0x1000])
|
||||
token_database.set_block_len([256])
|
||||
thread = worker.KVCacheStoreSendingThread(
|
||||
store=store,
|
||||
token_database=token_database,
|
||||
block_size=16,
|
||||
tp_rank=0,
|
||||
put_step=1,
|
||||
kv_role="kv_producer",
|
||||
ready_event=threading.Event(),
|
||||
)
|
||||
thread.request_queue.task_done = MagicMock()
|
||||
return thread
|
||||
|
||||
|
||||
def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:
|
||||
return ReqMeta(
|
||||
req_id=req_id,
|
||||
token_len_chunk=32,
|
||||
block_ids=[0, 1],
|
||||
block_hashes=block_hashes,
|
||||
can_save=True,
|
||||
original_block_size=16,
|
||||
)
|
||||
|
||||
|
||||
def test_store_sending_thread_skips_request_during_cpu_pressure():
|
||||
store = MagicMock()
|
||||
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
|
||||
store.batch_put_from_multi_buffers.side_effect = [
|
||||
[-200, -200],
|
||||
[256, 256],
|
||||
[256, 256],
|
||||
]
|
||||
thread = _make_store_sending_thread(store)
|
||||
|
||||
thread.add_stored_request("req-a")
|
||||
thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
|
||||
|
||||
assert thread._store_pressure_active is True
|
||||
assert "req-a" in thread._skip_store_requests
|
||||
assert store.batch_put_from_multi_buffers.call_count == 1
|
||||
|
||||
thread.add_stored_request("req-a")
|
||||
thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"]))
|
||||
|
||||
assert store.batch_put_from_multi_buffers.call_count == 1
|
||||
|
||||
thread.add_stored_request("req-b")
|
||||
thread._handle_request(_make_store_req("req-b", [b"b0", b"b1"]))
|
||||
|
||||
assert thread._store_pressure_active is False
|
||||
assert "req-a" not in thread._skip_store_requests
|
||||
assert store.batch_put_from_multi_buffers.call_count == 2
|
||||
|
||||
thread.add_stored_request("req-a")
|
||||
thread._handle_request(_make_store_req("req-a", [b"a4", b"a5"]))
|
||||
|
||||
assert store.batch_put_from_multi_buffers.call_count == 3
|
||||
|
||||
|
||||
def test_store_sending_thread_only_skips_on_no_available_handle():
|
||||
store = MagicMock()
|
||||
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
|
||||
store.batch_put_from_multi_buffers.side_effect = [
|
||||
[-500, -500],
|
||||
[256, 256],
|
||||
]
|
||||
thread = _make_store_sending_thread(store)
|
||||
|
||||
thread.add_stored_request("req-a")
|
||||
thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
|
||||
|
||||
assert thread._store_pressure_active is False
|
||||
assert "req-a" not in thread._skip_store_requests
|
||||
assert store.batch_put_from_multi_buffers.call_count == 1
|
||||
|
||||
thread.add_stored_request("req-a")
|
||||
thread._handle_request(_make_store_req("req-a", [b"a2", b"a3"]))
|
||||
|
||||
assert store.batch_put_from_multi_buffers.call_count == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers for register_kv_caches tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _auto_set_ready_event(*args, **kwargs):
|
||||
"""Side effect for mocked thread constructors that auto-sets ready_event."""
|
||||
for arg in args:
|
||||
if isinstance(arg, threading.Event):
|
||||
arg.set()
|
||||
for val in kwargs.values():
|
||||
if isinstance(val, threading.Event):
|
||||
val.set()
|
||||
return MagicMock()
|
||||
|
||||
|
||||
def _make_bare_worker(
|
||||
*,
|
||||
num_gpu_blocks: int = 10,
|
||||
block_size: int = 16,
|
||||
kv_role: str = "kv_both",
|
||||
) -> worker.MooncakeStoreWorker:
|
||||
"""Construct a MooncakeStoreWorker via __new__, bypassing __init__.
|
||||
|
||||
Sets only the attributes that register_kv_caches() reads so we can
|
||||
test the stride-based layout detection without a real
|
||||
MooncakeDistributedStore.
|
||||
"""
|
||||
w = object.__new__(worker.MooncakeStoreWorker)
|
||||
w.cache_config = MagicMock()
|
||||
w.cache_config.num_gpu_blocks = num_gpu_blocks
|
||||
w.store = MagicMock()
|
||||
w.store.register_buffer.return_value = 0
|
||||
w.use_mla = False
|
||||
w.token_database = ChunkedTokenDatabase(
|
||||
KeyMetadata("test-model", 0, 0, 0, 0), block_size=block_size
|
||||
)
|
||||
w.kv_role = kv_role
|
||||
w.block_size = block_size
|
||||
w.tp_rank = 0
|
||||
w.put_step = 1
|
||||
w.enable_kv_events = False
|
||||
w.kv_send_thread = None
|
||||
w.kv_recv_thread = None
|
||||
return w
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# register_kv_caches tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_register_kv_caches_blocks_first_single_segment():
|
||||
"""Blocks-first layout (FlashInfer/MLA): one segment per layer."""
|
||||
num_blocks = 10
|
||||
page_size_elements = 64 # elements per block
|
||||
w = _make_bare_worker(num_gpu_blocks=num_blocks)
|
||||
|
||||
# Shape: (num_blocks, page_size_elements) — blocks outermost, no outer_dims
|
||||
tensor = torch.zeros(num_blocks, page_size_elements, dtype=torch.float16)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"worker.KVCacheStoreSendingThread",
|
||||
side_effect=_auto_set_ready_event,
|
||||
),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"worker.KVCacheStoreRecvingThread",
|
||||
side_effect=_auto_set_ready_event,
|
||||
),
|
||||
):
|
||||
w.register_kv_caches({"layer0": tensor})
|
||||
|
||||
assert len(w.kv_caches_base_addr) == 1
|
||||
assert w.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
|
||||
|
||||
expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
|
||||
assert len(w.block_len) == 1
|
||||
assert w.block_len[0] == expected_block_len
|
||||
|
||||
w.store.register_buffer.assert_called_once_with(
|
||||
tensor.untyped_storage().data_ptr(),
|
||||
tensor.untyped_storage().nbytes(),
|
||||
)
|
||||
|
||||
|
||||
def test_register_kv_caches_kv_first_two_segments():
|
||||
"""K/V-first layout (FlashAttn): two segments (K, V) per layer."""
|
||||
num_blocks = 10
|
||||
block_size_tokens = 16
|
||||
num_kv_heads = 4
|
||||
head_size = 8
|
||||
|
||||
w = _make_bare_worker(num_gpu_blocks=num_blocks)
|
||||
|
||||
# Shape: (2, num_blocks, block_size, num_kv_heads, head_size) — K/V outermost
|
||||
tensor = torch.zeros(
|
||||
2,
|
||||
num_blocks,
|
||||
block_size_tokens,
|
||||
num_kv_heads,
|
||||
head_size,
|
||||
dtype=torch.float16,
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"worker.KVCacheStoreSendingThread",
|
||||
side_effect=_auto_set_ready_event,
|
||||
),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"worker.KVCacheStoreRecvingThread",
|
||||
side_effect=_auto_set_ready_event,
|
||||
),
|
||||
):
|
||||
w.register_kv_caches({"layer0": tensor})
|
||||
|
||||
# K/V-first: dim 0 has stride > page_size, so 2 segments
|
||||
assert len(w.kv_caches_base_addr) == 2
|
||||
assert len(w.block_len) == 2
|
||||
|
||||
el = tensor.element_size()
|
||||
seg_stride = tensor.stride(0) * el # stride of the K/V dim in bytes
|
||||
base = tensor.untyped_storage().data_ptr()
|
||||
assert w.kv_caches_base_addr[0] == base
|
||||
assert w.kv_caches_base_addr[1] == base + seg_stride
|
||||
assert w.block_len[0] == seg_stride // num_blocks
|
||||
assert w.block_len[1] == seg_stride // num_blocks
|
||||
|
||||
|
||||
def test_register_kv_caches_cross_layer_single_segment():
|
||||
"""Cross-layer tensor: single segment with block_len = page_size * num_layers."""
|
||||
num_blocks = 10
|
||||
num_layers = 4
|
||||
per_layer_page_elements = 64 # elements per layer per block
|
||||
|
||||
w = _make_bare_worker(num_gpu_blocks=num_blocks)
|
||||
|
||||
# Cross-layer blocks-first tensor: all layers packed into a single
|
||||
# contiguous block. Shape (num_blocks, num_layers * per_layer_page)
|
||||
# mimics the physical layout after stride reordering.
|
||||
total_page_elements = num_layers * per_layer_page_elements
|
||||
tensor = torch.zeros(num_blocks, total_page_elements, dtype=torch.float16)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"worker.KVCacheStoreSendingThread",
|
||||
side_effect=_auto_set_ready_event,
|
||||
),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"worker.KVCacheStoreRecvingThread",
|
||||
side_effect=_auto_set_ready_event,
|
||||
),
|
||||
):
|
||||
# Use the cross-layer wrapper key, same as register_cross_layers_kv_caches
|
||||
w.register_kv_caches({"__cross_layer__": tensor})
|
||||
|
||||
assert len(w.kv_caches_base_addr) == 1
|
||||
assert w.kv_caches_base_addr[0] == tensor.untyped_storage().data_ptr()
|
||||
|
||||
expected_block_len = tensor.untyped_storage().nbytes() // num_blocks
|
||||
# block_len should be per_layer_page_size * num_layers
|
||||
assert (
|
||||
expected_block_len
|
||||
== num_layers * per_layer_page_elements * tensor.element_size()
|
||||
)
|
||||
assert len(w.block_len) == 1
|
||||
assert w.block_len[0] == expected_block_len
|
||||
|
||||
# Also verify via register_cross_layers_kv_caches wrapper
|
||||
w2 = _make_bare_worker(num_gpu_blocks=num_blocks)
|
||||
with (
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"worker.KVCacheStoreSendingThread",
|
||||
side_effect=_auto_set_ready_event,
|
||||
),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
|
||||
"worker.KVCacheStoreRecvingThread",
|
||||
side_effect=_auto_set_ready_event,
|
||||
),
|
||||
):
|
||||
w2.register_cross_layers_kv_caches(tensor)
|
||||
|
||||
assert w2.kv_caches_base_addr == w.kv_caches_base_addr
|
||||
assert w2.block_len == w.block_len
|
||||
@@ -261,10 +261,11 @@ def test_multi_example_connector_consistency():
|
||||
)
|
||||
|
||||
events = get_connector_events()
|
||||
# First event is set_xfer_handshake_metadata from initialization, then
|
||||
# on_new_request when the request is enqueued, then
|
||||
# get_num_new_matched_tokens and update_state_after_alloc from generate().
|
||||
assert events["storage1-SCHEDULER"][:5] == [
|
||||
# First event is bind_gpu_block_pool from initialization, then
|
||||
# set_xfer_handshake_metadata, then on_new_request when the request is enqueued,
|
||||
# then get_num_new_matched_tokens and update_state_after_alloc from generate().
|
||||
assert events["storage1-SCHEDULER"][:6] == [
|
||||
"bind_gpu_block_pool",
|
||||
"set_xfer_handshake_metadata",
|
||||
"on_new_request",
|
||||
"get_num_new_matched_tokens 0",
|
||||
@@ -283,7 +284,8 @@ def test_multi_example_connector_consistency():
|
||||
"wait_for_layer_load",
|
||||
"save_kv_layer",
|
||||
]
|
||||
assert events["storage2-SCHEDULER"][:5] == [
|
||||
assert events["storage2-SCHEDULER"][:6] == [
|
||||
"bind_gpu_block_pool",
|
||||
"set_xfer_handshake_metadata",
|
||||
"on_new_request",
|
||||
"get_num_new_matched_tokens 0",
|
||||
|
||||
@@ -93,27 +93,39 @@ def test_logical_to_kernel_block_ids_with_hma():
|
||||
|
||||
@pytest.mark.cpu_test
|
||||
@pytest.mark.parametrize(
|
||||
"group_spec_types,expansion_stride,remote_block_ids,expected_remote_block_ids",
|
||||
"group_spec_types,remote_physical_per_logical,"
|
||||
"local_physical_per_logical,tp_ratio,remote_block_ids,"
|
||||
"expected_remote_block_ids",
|
||||
[
|
||||
pytest.param(
|
||||
("FullAttentionSpec", "SlidingWindowSpec"),
|
||||
2,
|
||||
2,
|
||||
1,
|
||||
([0, 1, 2], [3, 4]),
|
||||
[[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]],
|
||||
id="dense_fa_swa",
|
||||
),
|
||||
# Nemotron-3-Nano-30B-A3B 4p1d (P_TP=4, D_TP=1):
|
||||
# remote_physical_per_logical=34, local_physical_per_logical=66.
|
||||
# FA logical block 5 → kernel [170..203], block 6 → [204..237].
|
||||
# Mamba block unchanged.
|
||||
pytest.param(
|
||||
("FullAttentionSpec", "MambaSpec"),
|
||||
261,
|
||||
([0, 1, 2], [10, 11]),
|
||||
[[0, 1, 261, 262, 522, 523], [10, 11]],
|
||||
34,
|
||||
66,
|
||||
-4,
|
||||
([5, 6], [2]),
|
||||
[list(range(170, 238)), [2]],
|
||||
id="mamba_fa_ssm",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_read_blocks_for_req_expands_remote_ids(
|
||||
group_spec_types,
|
||||
expansion_stride,
|
||||
remote_physical_per_logical,
|
||||
local_physical_per_logical,
|
||||
tp_ratio,
|
||||
remote_block_ids,
|
||||
expected_remote_block_ids,
|
||||
):
|
||||
@@ -148,7 +160,7 @@ def test_read_blocks_for_req_expands_remote_ids(
|
||||
resolved_types = tuple(spec_name_to_type[n] for n in group_spec_types)
|
||||
|
||||
worker = object.__new__(NixlConnectorWorker)
|
||||
worker._physical_blocks_per_logical_kv_block = 2
|
||||
worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
|
||||
|
||||
has_mamba = any(t is MambaSpec for t in resolved_types)
|
||||
has_swa = any(t is SlidingWindowSpec for t in resolved_types)
|
||||
@@ -159,9 +171,11 @@ def test_read_blocks_for_req_expands_remote_ids(
|
||||
remote_engine_id = "remote-engine"
|
||||
|
||||
worker.transfer_topo = MagicMock()
|
||||
worker.transfer_topo.tp_ratio.return_value = 1
|
||||
# tp_ratio not exercised (all_source_ranks is empty so no reads run),
|
||||
# but set for realism.
|
||||
worker.transfer_topo.tp_ratio.return_value = tp_ratio
|
||||
remote_info = MagicMock()
|
||||
remote_info.remote_physical_blocks_per_logical = expansion_stride
|
||||
remote_info.remote_physical_blocks_per_logical = remote_physical_per_logical
|
||||
worker.transfer_topo.get_engine_info.return_value = remote_info
|
||||
worker.use_mla = False
|
||||
|
||||
@@ -192,6 +206,168 @@ def test_read_blocks_for_req_expands_remote_ids(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.cpu_test
|
||||
@pytest.mark.parametrize(
|
||||
"local_physical_per_logical,remote_physical_per_logical,"
|
||||
"local_block_ids,remote_block_ids,"
|
||||
"expected_local,expected_remote",
|
||||
[
|
||||
# 10 kernel blocks of data, local has more logical blocks.
|
||||
# remote physical_per_logical=10 → 1 logical → 10 kernel blocks
|
||||
# local physical_per_logical=6 → 2 logical → 12 kernel blocks
|
||||
# Trim local from 12 to 10.
|
||||
pytest.param(
|
||||
6,
|
||||
10,
|
||||
[list(range(12)), [42]],
|
||||
[list(range(10)), [42]],
|
||||
[list(range(10)), [42]],
|
||||
[list(range(10)), [42]],
|
||||
id="align_local6_remote10",
|
||||
),
|
||||
# 10 kernel blocks of data, remote has more logical blocks.
|
||||
# remote physical_per_logical=6 → 2 logical → 12 kernel blocks
|
||||
# local physical_per_logical=10 → 1 logical → 10 kernel blocks
|
||||
# Trim remote from 12 to 10.
|
||||
pytest.param(
|
||||
10,
|
||||
6,
|
||||
[list(range(10)), [42]],
|
||||
[list(range(12)), [42]],
|
||||
[list(range(10)), [42]],
|
||||
[list(range(10)), [42]],
|
||||
id="align_local10_remote6",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_apply_prefix_caching_mamba_hybrid(
|
||||
local_physical_per_logical,
|
||||
remote_physical_per_logical,
|
||||
local_block_ids,
|
||||
remote_block_ids,
|
||||
expected_local,
|
||||
expected_remote,
|
||||
):
|
||||
"""_apply_prefix_caching front-trims FA groups to
|
||||
min(local, remote) for Mamba hybrid models with heterogeneous TP.
|
||||
"""
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
|
||||
NixlConnectorWorker,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec
|
||||
|
||||
worker = object.__new__(NixlConnectorWorker)
|
||||
worker._has_mamba = True
|
||||
worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
|
||||
worker._group_spec_types = (FullAttentionSpec, MambaSpec)
|
||||
worker.kv_cache_config = make_kv_cache_config(block_size=16, mamba_enabled=True)
|
||||
|
||||
aligned_local, aligned_remote = worker._apply_prefix_caching(
|
||||
local_block_ids, remote_block_ids, remote_physical_per_logical
|
||||
)
|
||||
|
||||
assert aligned_local == expected_local, (
|
||||
f"Expected local {expected_local}, got {aligned_local}"
|
||||
)
|
||||
assert aligned_remote == expected_remote, (
|
||||
f"Expected remote {expected_remote}, got {aligned_remote}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.cpu_test
|
||||
@pytest.mark.parametrize(
|
||||
"local_physical_per_logical,remote_physical_per_logical,"
|
||||
"remote_fa_blocks,local_fa_blocks,ssm_blocks,"
|
||||
"correct_remote_fa,correct_local_fa",
|
||||
[
|
||||
# 10 kernel blocks of data (640 tokens).
|
||||
# remote physical_per_logical=10 → 1 logical → 10 kernel [0..9]
|
||||
# local physical_per_logical=6 → 2 logical → 12 kernel [0..11]
|
||||
# 1st local logical block cached → suffix [6..11]
|
||||
# Correct: transfer only uncached suffix tokens (384-639)
|
||||
# = remote [6,7,8,9] → local [6,7,8,9].
|
||||
# Actual (front-trim): remote[:6]=[0..5] → local [6..11]. Wrong.
|
||||
pytest.param(
|
||||
6,
|
||||
10,
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
||||
[6, 7, 8, 9, 10, 11],
|
||||
[42],
|
||||
[6, 7, 8, 9],
|
||||
[6, 7, 8, 9],
|
||||
id="local6_remote10_fail",
|
||||
),
|
||||
# 15 kernel blocks of data (960 tokens).
|
||||
# remote physical_per_logical=6 → 3 logical → 18 kernel [0..17]
|
||||
# local physical_per_logical=10 → 2 logical → 20 kernel [0..19]
|
||||
# 1st local logical block cached → suffix [10..19]
|
||||
# Correct: transfer only uncached suffix tokens (640-959)
|
||||
# = remote [10,11,12,13,14] → local [10,11,12,13,14].
|
||||
# Actual (front-trim): remote[:10]=[0..9] → local [10..19]. Wrong.
|
||||
pytest.param(
|
||||
10,
|
||||
6,
|
||||
list(range(18)),
|
||||
list(range(10, 20)),
|
||||
[42],
|
||||
[10, 11, 12, 13, 14],
|
||||
[10, 11, 12, 13, 14],
|
||||
id="local10_remote6_fail",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_mismatched_physical_per_logical_fails_with_prefix_caching(
|
||||
local_physical_per_logical,
|
||||
remote_physical_per_logical,
|
||||
remote_fa_blocks,
|
||||
local_fa_blocks,
|
||||
ssm_blocks,
|
||||
correct_remote_fa,
|
||||
correct_local_fa,
|
||||
):
|
||||
"""Demonstrate that _apply_prefix_caching front-trims ([:N])
|
||||
in the Mamba hybrid path, which fails when prefix caching produces
|
||||
suffix-only local blocks.
|
||||
|
||||
Prefix caching operates at logical block granularity. When a logical
|
||||
block is cached locally, the decode side only allocates kernel blocks
|
||||
for the uncached suffix. The front-trim pairs remote prefix blocks
|
||||
with local suffix slots — a silent data corruption.
|
||||
"""
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
|
||||
NixlConnectorWorker,
|
||||
)
|
||||
|
||||
worker = object.__new__(NixlConnectorWorker)
|
||||
worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
|
||||
worker.kv_cache_config = make_kv_cache_config(
|
||||
block_size=16,
|
||||
mamba_enabled=True,
|
||||
)
|
||||
worker._has_mamba = True
|
||||
worker._group_spec_types = tuple(
|
||||
type(g.kv_cache_spec) for g in worker.kv_cache_config.kv_cache_groups
|
||||
)
|
||||
|
||||
local_block_ids = (local_fa_blocks, ssm_blocks)
|
||||
remote_block_ids = (remote_fa_blocks, ssm_blocks)
|
||||
|
||||
aligned_local, aligned_remote = worker._apply_prefix_caching(
|
||||
local_block_ids,
|
||||
remote_block_ids,
|
||||
remote_physical_per_logical,
|
||||
)
|
||||
|
||||
assert (
|
||||
aligned_remote[0] != correct_remote_fa or aligned_local[0] != correct_local_fa
|
||||
), (
|
||||
f"Prefix caching with mismatched physical_per_logical should not "
|
||||
f"produce correct transfer ids: "
|
||||
f"remote={aligned_remote[0]}, local={aligned_local[0]}, "
|
||||
f"correct_remote={correct_remote_fa}, correct_local={correct_local_fa}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
|
||||
def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size):
|
||||
"""Test that a prefill instance returns fewer "remote blocks" for the SWA groups
|
||||
@@ -564,3 +740,119 @@ def test_compute_physical_blocks_per_logical(ssm_sizes, block_len, expected_rati
|
||||
)
|
||||
|
||||
assert compute_physical_blocks_per_logical(ssm_sizes, block_len) == expected_ratio
|
||||
|
||||
|
||||
@pytest.mark.cpu_test
|
||||
@pytest.mark.parametrize(
|
||||
"mamba_enabled,swa_enabled,"
|
||||
"local_physical_per_logical,remote_physical_per_logical,"
|
||||
"logical_block_ids,expected_kernel_block_ids",
|
||||
[
|
||||
# Qwen3.5-0.8B 4P2D (kernel_block_size=64):
|
||||
# prefill TP=4: logical_block_size=384 → physical_per_logical=6
|
||||
# decode TP=2: logical_block_size=640 → physical_per_logical=10
|
||||
# FA logical [0] → remote kernel [0..9] (1 * 10)
|
||||
# SSM logical [10] → unchanged [10]
|
||||
pytest.param(
|
||||
True,
|
||||
False,
|
||||
6,
|
||||
10,
|
||||
([0], [10]),
|
||||
[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10]],
|
||||
id="qwen35_4p2d",
|
||||
),
|
||||
# Qwen3.5-0.8B 2P4D (kernel_block_size=64):
|
||||
# prefill TP=2: logical_block_size=640 → physical_per_logical=10
|
||||
# decode TP=4: logical_block_size=384 → physical_per_logical=6
|
||||
# FA logical [0, 1] → remote kernel [0..5, 6..11] (2 * 6)
|
||||
# SSM logical [10] → unchanged [10]
|
||||
pytest.param(
|
||||
True,
|
||||
False,
|
||||
10,
|
||||
6,
|
||||
([0, 1], [10]),
|
||||
[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], [10]],
|
||||
id="qwen35_2p4d",
|
||||
),
|
||||
# Homogeneous TP (kernel_block_size=64):
|
||||
# both sides: logical_block_size=640 → physical_per_logical=10
|
||||
# FA logical [0] → kernel [0..9], SSM unchanged
|
||||
pytest.param(
|
||||
True,
|
||||
False,
|
||||
10,
|
||||
10,
|
||||
([0], [10]),
|
||||
[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10]],
|
||||
id="homo_tp",
|
||||
),
|
||||
# remote physical_per_logical=1: early return, no expansion
|
||||
pytest.param(
|
||||
True,
|
||||
False,
|
||||
10,
|
||||
1,
|
||||
([0, 1, 2], [5]),
|
||||
[[0, 1, 2], [5]],
|
||||
id="mamba_remote_physical_per_logical_1",
|
||||
),
|
||||
# Pure FA (no mamba): single group expanded with remote stride
|
||||
pytest.param(
|
||||
False,
|
||||
False,
|
||||
2,
|
||||
4,
|
||||
([0, 1],),
|
||||
[[0, 1, 2, 3, 4, 5, 6, 7]],
|
||||
id="pure_fa",
|
||||
),
|
||||
# FA + SWA (no mamba): both groups expanded
|
||||
pytest.param(
|
||||
False,
|
||||
True,
|
||||
2,
|
||||
3,
|
||||
([0, 1], [2, 3]),
|
||||
[[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]],
|
||||
id="fa_swa",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_logical_to_remote_kernel_block_ids(
|
||||
mamba_enabled,
|
||||
swa_enabled,
|
||||
local_physical_per_logical,
|
||||
remote_physical_per_logical,
|
||||
logical_block_ids,
|
||||
expected_kernel_block_ids,
|
||||
):
|
||||
"""Verify _logical_to_remote_kernel_block_ids uses the remote
|
||||
physical_per_logical for FA expansion, not the local one.
|
||||
|
||||
This was the root cause of silent accuracy corruption in Qwen3.5
|
||||
heterogeneous TP (e.g. 4P2D): the old code used local physical_per_logical
|
||||
for the expansion arange, producing wrong kernel block indices.
|
||||
|
||||
Qwen3.5-0.8B values verified by verify_conv_split.py (issue #13).
|
||||
"""
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
|
||||
NixlConnectorWorker,
|
||||
)
|
||||
|
||||
worker = object.__new__(NixlConnectorWorker)
|
||||
worker._physical_blocks_per_logical_kv_block = local_physical_per_logical
|
||||
worker.kv_cache_config = make_kv_cache_config(
|
||||
block_size=16,
|
||||
mamba_enabled=mamba_enabled,
|
||||
swa_enabled=swa_enabled,
|
||||
)
|
||||
|
||||
result = worker._logical_to_remote_kernel_block_ids(
|
||||
logical_block_ids,
|
||||
remote_physical_per_logical,
|
||||
)
|
||||
assert list(result) == expected_kernel_block_ids, (
|
||||
f"Expected {expected_kernel_block_ids}, got {result}"
|
||||
)
|
||||
|
||||
@@ -19,9 +19,11 @@ from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
|
||||
from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy
|
||||
|
||||
|
||||
def make_req_context(kv_transfer_params: dict | None = None) -> ReqContext:
|
||||
def make_req_context(
|
||||
req_id: str = "", kv_transfer_params: dict | None = None
|
||||
) -> ReqContext:
|
||||
"""Create a ReqContext as production code would, from a request's params."""
|
||||
return ReqContext(kv_transfer_params=kv_transfer_params)
|
||||
return ReqContext(req_id=req_id, kv_transfer_params=kv_transfer_params)
|
||||
|
||||
|
||||
_EMPTY_REQ_CTX = make_req_context()
|
||||
|
||||
@@ -179,6 +179,7 @@ _req_counter = 0
|
||||
def make_request(
|
||||
num_blocks: int = 2,
|
||||
request_id: str | None = None,
|
||||
extra_tokens: int = 1,
|
||||
) -> Request:
|
||||
"""Create a Request with deterministic block hashes."""
|
||||
global _req_counter
|
||||
@@ -186,13 +187,7 @@ def make_request(
|
||||
if request_id is None:
|
||||
request_id = f"req-{_req_counter}"
|
||||
|
||||
# Add one extra token beyond the last full block so that
|
||||
# ``max_cache_hit_length = num_tokens - 1`` (see
|
||||
# KVCacheManager.get_computed_blocks) does not truncate the final
|
||||
# full block: ``find_longest_cache_hit`` uses
|
||||
# ``max_length // block_size`` and would otherwise drop one block
|
||||
# when the prompt is an exact multiple of block_size.
|
||||
num_tokens = num_blocks * BLOCK_SIZE + 1
|
||||
num_tokens = num_blocks * BLOCK_SIZE + extra_tokens
|
||||
start = _req_counter * 10000
|
||||
prompt_token_ids = list(range(start, start + num_tokens))
|
||||
sampling_params = SamplingParams(max_tokens=1)
|
||||
@@ -386,6 +381,8 @@ def test_eager_store_and_load_roundtrip() -> None:
|
||||
block_hasher=req._block_hasher,
|
||||
)
|
||||
hit_tokens, is_async = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
|
||||
# make_request pads num_tokens by +1 beyond the last full block, so the
|
||||
# manager's max_hit_len = num_tokens - 1 cap leaves all full blocks intact.
|
||||
assert hit_tokens == num_blocks * BLOCK_SIZE
|
||||
assert is_async is True
|
||||
|
||||
@@ -405,7 +402,44 @@ def test_eager_store_and_load_roundtrip() -> None:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1b: Lazy store-and-load roundtrip
|
||||
# Test 1b: Boundary — max_hit_len cap drops the last full block when the
|
||||
# prompt is an exact multiple of BLOCK_SIZE.
|
||||
# ---------------------------------------------------------------------------
|
||||
def test_max_hit_len_cap_drops_last_full_block() -> None:
|
||||
"""When num_tokens is an exact multiple of BLOCK_SIZE, the manager's
|
||||
``max_hit_len = num_tokens - 1`` cap forces ``find_longest_cache_hit`` to
|
||||
drop the final block (since ``max_length // block_size`` rounds down).
|
||||
"""
|
||||
fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
|
||||
sched = fix.scheduler
|
||||
|
||||
num_blocks = 2
|
||||
req = make_request(num_blocks=num_blocks, extra_tokens=0)
|
||||
assert req.num_tokens == num_blocks * BLOCK_SIZE
|
||||
|
||||
kv_blocks = _alloc_and_register(fix, req, num_blocks)
|
||||
sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
|
||||
sched_out = make_scheduler_output(
|
||||
{req.request_id: num_blocks * BLOCK_SIZE},
|
||||
new_reqs={req.request_id: kv_blocks.get_block_ids()},
|
||||
)
|
||||
meta = sched.build_connector_meta(sched_out)
|
||||
simulate_store_completion(sched, meta.store_event)
|
||||
|
||||
req2 = Request(
|
||||
request_id="req-cap-boundary",
|
||||
prompt_token_ids=req.prompt_token_ids,
|
||||
sampling_params=req.sampling_params,
|
||||
pooling_params=None,
|
||||
mm_features=None,
|
||||
block_hasher=req._block_hasher,
|
||||
)
|
||||
hit_tokens, _ = sched.get_num_new_matched_tokens(req2, num_computed_tokens=0)
|
||||
assert hit_tokens == (num_blocks - 1) * BLOCK_SIZE
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1c: Lazy store-and-load roundtrip
|
||||
# ---------------------------------------------------------------------------
|
||||
def _flush_old_blocks_to_lru_head(
|
||||
gpu_pool: BlockPool,
|
||||
@@ -469,8 +503,11 @@ def test_lazy_store_and_load_roundtrip() -> None:
|
||||
hit_tokens, is_async = sched.get_num_new_matched_tokens(
|
||||
req_old2, num_computed_tokens=0
|
||||
)
|
||||
assert hit_tokens == num_blocks * BLOCK_SIZE, (
|
||||
f"Expected {num_blocks * BLOCK_SIZE} hit tokens, got {hit_tokens}"
|
||||
# make_request pads num_tokens by +1 beyond the last full block, so the
|
||||
# manager's max_hit_len = num_tokens - 1 cap leaves all full blocks intact.
|
||||
expected_hit = num_blocks * BLOCK_SIZE
|
||||
assert hit_tokens == expected_hit, (
|
||||
f"Expected {expected_hit} hit tokens, got {hit_tokens}"
|
||||
)
|
||||
assert is_async is True
|
||||
|
||||
@@ -537,7 +574,72 @@ def test_eager_duplicate_store_skipped() -> None:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2b: Lazy duplicate store is skipped
|
||||
# Test 2b: Eager dedup of in-flight stores across consecutive steps
|
||||
# ---------------------------------------------------------------------------
|
||||
def test_eager_in_flight_store_dedup_across_steps() -> None:
|
||||
"""Eager: a second request sharing a prefix with an in-flight store
|
||||
must not re-offload the same GPU blocks before completion lands.
|
||||
|
||||
Simulates a GPU prefix-cache hit by reusing the first request's
|
||||
GPU block IDs in the second scheduler step, which is the path the
|
||||
real scheduler takes when two requests share a prefix.
|
||||
"""
|
||||
fix = make_scheduler(num_cpu_blocks=8, num_gpu_blocks=16, lazy=False)
|
||||
sched = fix.scheduler
|
||||
|
||||
num_blocks = 2
|
||||
req = make_request(num_blocks=num_blocks)
|
||||
|
||||
kv_blocks = _alloc_and_register(fix, req, num_blocks)
|
||||
sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
|
||||
block_ids = kv_blocks.get_block_ids()
|
||||
sched_out = make_scheduler_output(
|
||||
{req.request_id: num_blocks * BLOCK_SIZE},
|
||||
new_reqs={req.request_id: block_ids},
|
||||
)
|
||||
|
||||
meta1 = sched.build_connector_meta(sched_out)
|
||||
assert meta1.store_event >= 0
|
||||
assert len(meta1.store_cpu_blocks) == num_blocks
|
||||
# In-flight set tracks the scheduled GPU blocks until completion.
|
||||
assert sched._in_flight_store_gpu_blocks == set(meta1.store_gpu_blocks)
|
||||
cpu_free_after_first = get_cpu_free_blocks(sched)
|
||||
|
||||
# Second request shares the prefix and reuses the same GPU block IDs
|
||||
# (the real scheduler path: GPU prefix cache returns the same blocks).
|
||||
# Do NOT simulate completion — the first store is still in-flight.
|
||||
req2 = Request(
|
||||
request_id="req-dup-eager-inflight",
|
||||
prompt_token_ids=req.prompt_token_ids,
|
||||
sampling_params=req.sampling_params,
|
||||
pooling_params=None,
|
||||
mm_features=None,
|
||||
block_hasher=req._block_hasher,
|
||||
)
|
||||
req2.num_computed_tokens = num_blocks * BLOCK_SIZE
|
||||
sched.update_state_after_alloc(req2, kv_blocks, num_external_tokens=0)
|
||||
sched_out2 = make_scheduler_output(
|
||||
{req2.request_id: num_blocks * BLOCK_SIZE},
|
||||
new_reqs={req2.request_id: block_ids},
|
||||
)
|
||||
|
||||
meta2 = sched.build_connector_meta(sched_out2)
|
||||
if meta2.store_event >= 0:
|
||||
assert len(meta2.store_cpu_blocks) == 0, (
|
||||
"Expected no new CPU blocks for in-flight duplicate hashes"
|
||||
)
|
||||
assert get_cpu_free_blocks(sched) == cpu_free_after_first, (
|
||||
"Second request should not consume CPU blocks while the first "
|
||||
"store is still in-flight"
|
||||
)
|
||||
|
||||
# After completion, the in-flight set is cleared.
|
||||
simulate_store_completion(sched, meta1.store_event)
|
||||
assert sched._in_flight_store_gpu_blocks == set()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2c: Lazy duplicate store is skipped
|
||||
# ---------------------------------------------------------------------------
|
||||
def test_lazy_duplicate_store_skipped() -> None:
|
||||
"""Lazy: blocks already offloaded to CPU should not be offloaded again.
|
||||
@@ -1092,11 +1194,9 @@ def test_partial_gpu_prefix_plus_cpu_load() -> None:
|
||||
hit_tokens, is_async = sched.get_num_new_matched_tokens(
|
||||
req2, num_computed_tokens=gpu_local_computed
|
||||
)
|
||||
# CPU should hit blocks 2,3 (not 4,5 — those are beyond the CPU range).
|
||||
num_cpu_hit_blocks = 2
|
||||
# Actually CPU has all 6 stored; it returns hits starting from position 2.
|
||||
# The number of CPU hit blocks = min(remaining request blocks, CPU cached).
|
||||
# Here remaining = 6 - 2 = 4 blocks are in CPU, so hit = 4 * BLOCK_SIZE.
|
||||
# CPU has all 6 blocks stored. make_request pads num_tokens by +1, so
|
||||
# the manager's num_tokens - 1 cap leaves all full blocks intact:
|
||||
# remaining hashable range = 6 - 2 = 4 blocks, all hit.
|
||||
num_cpu_hit_blocks = 4
|
||||
assert hit_tokens == num_cpu_hit_blocks * BLOCK_SIZE, (
|
||||
f"Expected {num_cpu_hit_blocks * BLOCK_SIZE} CPU hit tokens, got {hit_tokens}"
|
||||
|
||||
+71
-4
@@ -6,8 +6,8 @@ import math
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.v1.worker.gpu.spec_decode.probabilistic_rejection_sampler_utils import (
|
||||
probabilistic_rejection_sample,
|
||||
from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
|
||||
rejection_sample,
|
||||
)
|
||||
|
||||
VOCAB_SIZE = 4096
|
||||
@@ -167,7 +167,7 @@ def test_stochastic_rejection_sample(num_speculative_steps: int, temperature: fl
|
||||
num_trials=num_trials,
|
||||
)
|
||||
|
||||
sampled, num_sampled = probabilistic_rejection_sample(
|
||||
sampled, num_sampled = rejection_sample(
|
||||
**inputs, num_speculative_steps=num_speculative_steps
|
||||
)
|
||||
|
||||
@@ -201,7 +201,7 @@ def test_greedy_rejection_sample(num_speculative_steps: int):
|
||||
num_trials=num_trials,
|
||||
)
|
||||
|
||||
sampled, num_sampled = probabilistic_rejection_sample(
|
||||
sampled, num_sampled = rejection_sample(
|
||||
**inputs, num_speculative_steps=num_speculative_steps
|
||||
)
|
||||
|
||||
@@ -213,3 +213,70 @@ def test_greedy_rejection_sample(num_speculative_steps: int):
|
||||
assert (sampled[accepted_mask] == target_argmax).all(), (
|
||||
"Greedy sampling produced tokens that are not the target argmax"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_speculative_steps,temperature,unconditional_rates",
|
||||
[
|
||||
(3, 1.0, [0.9, 0.5, 0.2]),
|
||||
(3, 0.0, [0.9, 0.5, 0.2]),
|
||||
(3, 1.0, [1.0, 1.0, 1.0]),
|
||||
(3, 0.0, [1.0, 1.0, 1.0]),
|
||||
(3, 1.0, [0.0, 0.0, 0.0]),
|
||||
(3, 0.0, [0.0, 0.0, 0.0]),
|
||||
(1, 1.0, [0.7]),
|
||||
(1, 0.0, [0.7]),
|
||||
],
|
||||
)
|
||||
def test_synthetic_rejection_sample(
|
||||
num_speculative_steps: int,
|
||||
temperature: float,
|
||||
unconditional_rates: list[float],
|
||||
):
|
||||
"""
|
||||
Verify that synthetic rejection sampling produces the expected
|
||||
per-position acceptance rates. The unconditional rate at position i
|
||||
is P(all draft steps 0..i accepted) = product(conditional_rates[0:i+1]).
|
||||
This is approximately mean(num accepted >= i + 1) over many trials.
|
||||
"""
|
||||
from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates
|
||||
|
||||
torch.manual_seed(42)
|
||||
device = "cuda"
|
||||
num_trials = 10 * VOCAB_SIZE
|
||||
deviation_tol = 1e-2
|
||||
|
||||
target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
|
||||
draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
|
||||
|
||||
if temperature > 0:
|
||||
target_logits_1d /= temperature
|
||||
draft_logits_1d /= temperature
|
||||
|
||||
inputs = _build_rejection_sample_inputs(
|
||||
target_logits_1d,
|
||||
draft_logits_1d,
|
||||
num_speculative_steps,
|
||||
temperature=temperature,
|
||||
num_trials=num_trials,
|
||||
)
|
||||
|
||||
conditional_rates = unconditional_to_conditional_rates(unconditional_rates)
|
||||
synthetic_conditional_rates = torch.tensor(
|
||||
conditional_rates, dtype=torch.float32, device=device
|
||||
)
|
||||
|
||||
_, num_sampled = rejection_sample(
|
||||
**inputs,
|
||||
num_speculative_steps=num_speculative_steps,
|
||||
synthetic_conditional_rates=synthetic_conditional_rates,
|
||||
)
|
||||
|
||||
# num_sampled includes the resampled/bonus token.
|
||||
num_accepted = num_sampled - 1
|
||||
for i, expected_rate in enumerate(unconditional_rates):
|
||||
observed_rate = (num_accepted >= i + 1).float().mean().item()
|
||||
assert abs(observed_rate - expected_rate) < deviation_tol, (
|
||||
f"Step {i}: observed rate {observed_rate:.4f} deviates from "
|
||||
f"expected rate {expected_rate:.4f} by more than {deviation_tol}."
|
||||
)
|
||||
+82
-44
@@ -1,5 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import dataclasses
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -9,22 +11,55 @@ from vllm import LLM
|
||||
from vllm.config import SpeculativeConfig
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
|
||||
MODEL_PATH = "nm-testing/dflash-qwen3-8b-speculators"
|
||||
|
||||
EXPECTED_GSM8K_ACCURACY = 0.885
|
||||
ACCURACY_RTOL = 0.03
|
||||
EXPECTED_ACCEPTANCE_LEN = 3.45
|
||||
ACCEPTANCE_LEN_RTOL = 0.15
|
||||
|
||||
# Expected per-position acceptance rates (accepted_at_pos / num_drafts)
|
||||
# Based on GSM8K evaluation with Qwen3-8B dflash speculators.
|
||||
EXPECTED_PER_POS_ACCEPTANCE_RATES = [0.795, 0.611, 0.429, 0.282]
|
||||
PER_POS_RTOL = 0.15
|
||||
@dataclasses.dataclass
|
||||
class SpeculatorTestConfig:
|
||||
model_path: str
|
||||
method: str
|
||||
display_name: str
|
||||
expected_gsm8k_accuracy: float
|
||||
accuracy_rtol: float
|
||||
expected_acceptance_len: float
|
||||
acceptance_len_rtol: float
|
||||
expected_per_pos_acceptance_rates: tuple[float, ...]
|
||||
per_pos_rtol: float
|
||||
quantization: str | None = None
|
||||
parallel_drafting: bool | None = None
|
||||
|
||||
|
||||
def compute_spec_decode_stats(
|
||||
metrics,
|
||||
) -> dict:
|
||||
DFLASH_CONFIG = SpeculatorTestConfig(
|
||||
model_path="nm-testing/dflash-qwen3-8b-speculators",
|
||||
method="dflash",
|
||||
display_name="DFlash",
|
||||
expected_gsm8k_accuracy=0.885,
|
||||
accuracy_rtol=0.03,
|
||||
expected_acceptance_len=3.45,
|
||||
acceptance_len_rtol=0.15,
|
||||
expected_per_pos_acceptance_rates=(0.795, 0.611, 0.429, 0.282),
|
||||
per_pos_rtol=0.15,
|
||||
quantization="fp8",
|
||||
)
|
||||
|
||||
PEAGLE_CONFIG = SpeculatorTestConfig(
|
||||
model_path="nm-testing/qwen3-8b-peagle-speculators",
|
||||
method="eagle3",
|
||||
display_name="PEagle",
|
||||
expected_gsm8k_accuracy=0.88,
|
||||
accuracy_rtol=0.05,
|
||||
expected_acceptance_len=2.27,
|
||||
acceptance_len_rtol=0.20,
|
||||
expected_per_pos_acceptance_rates=(0.66, 0.36, 0.18, 0.09),
|
||||
per_pos_rtol=0.20,
|
||||
parallel_drafting=True,
|
||||
)
|
||||
|
||||
SPECULATOR_CONFIGS = [
|
||||
pytest.param(DFLASH_CONFIG, id="dflash"),
|
||||
pytest.param(PEAGLE_CONFIG, id="peagle"),
|
||||
]
|
||||
|
||||
|
||||
def compute_spec_decode_stats(metrics) -> dict:
|
||||
"""Extract all spec-decode metrics and compute derived stats."""
|
||||
name2metric = {m.name: m for m in metrics}
|
||||
|
||||
@@ -67,25 +102,26 @@ def print_spec_decode_stats(stats: dict) -> None:
|
||||
print("===============================\n")
|
||||
|
||||
|
||||
def test_dflash_speculators_model(vllm_runner, example_prompts, monkeypatch):
|
||||
@pytest.mark.parametrize("config", SPECULATOR_CONFIGS)
|
||||
def test_speculators_model(vllm_runner, example_prompts, monkeypatch, config):
|
||||
"""
|
||||
Test DFlash speculators model properly initializes speculative decoding.
|
||||
Test speculators model properly initializes speculative decoding.
|
||||
|
||||
Verifies:
|
||||
1. Speculative config is automatically initialized from speculators config
|
||||
2. Method is detected as 'dflash'
|
||||
3. The draft model path is correctly set
|
||||
4. Speculative tokens count is valid (num_speculative_tokens=8)
|
||||
5. Text generation works with speculative decoding enabled
|
||||
2. Method is detected correctly
|
||||
3. parallel_drafting is set correctly (if applicable)
|
||||
4. The draft model path is correctly set
|
||||
5. Speculative tokens count is valid
|
||||
6. Text generation works with speculative decoding enabled
|
||||
"""
|
||||
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
|
||||
with vllm_runner(
|
||||
MODEL_PATH,
|
||||
dtype=torch.bfloat16,
|
||||
enforce_eager=True,
|
||||
quantization="fp8",
|
||||
) as vllm_model:
|
||||
runner_kwargs = dict(dtype=torch.bfloat16, enforce_eager=True)
|
||||
if config.quantization:
|
||||
runner_kwargs["quantization"] = config.quantization
|
||||
|
||||
with vllm_runner(config.model_path, **runner_kwargs) as vllm_model:
|
||||
vllm_config = vllm_model.llm.llm_engine.vllm_config
|
||||
|
||||
assert isinstance(vllm_config.speculative_config, SpeculativeConfig), (
|
||||
@@ -93,40 +129,43 @@ def test_dflash_speculators_model(vllm_runner, example_prompts, monkeypatch):
|
||||
)
|
||||
|
||||
spec_config = vllm_config.speculative_config
|
||||
assert spec_config.method == "dflash", (
|
||||
f"Expected method='dflash', got '{spec_config.method}'"
|
||||
assert spec_config.method == config.method, (
|
||||
f"Expected method='{config.method}', got '{spec_config.method}'"
|
||||
)
|
||||
if config.parallel_drafting is not None:
|
||||
assert spec_config.parallel_drafting is config.parallel_drafting, (
|
||||
f"Expected parallel_drafting={config.parallel_drafting} "
|
||||
f"for {config.display_name} model"
|
||||
)
|
||||
assert spec_config.num_speculative_tokens > 0, (
|
||||
f"Expected positive speculative tokens, "
|
||||
f"got {spec_config.num_speculative_tokens}"
|
||||
)
|
||||
assert spec_config.model == MODEL_PATH, (
|
||||
f"Draft model should be {MODEL_PATH}, got {spec_config.model}"
|
||||
assert spec_config.model == config.model_path, (
|
||||
f"Draft model should be {config.model_path}, got {spec_config.model}"
|
||||
)
|
||||
|
||||
vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens=20)
|
||||
assert vllm_outputs, f"No outputs generated for speculators model {MODEL_PATH}"
|
||||
assert vllm_outputs, (
|
||||
f"No outputs generated for speculators model {config.model_path}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow_test
|
||||
@large_gpu_mark(min_gb=40)
|
||||
def test_dflash_speculators_correctness(monkeypatch):
|
||||
@pytest.mark.parametrize("config", SPECULATOR_CONFIGS)
|
||||
def test_speculators_correctness(monkeypatch, config):
|
||||
"""
|
||||
E2E correctness test for DFlash via the speculators auto-detect path.
|
||||
E2E correctness test via the speculators auto-detect path.
|
||||
|
||||
Evaluates GSM8k accuracy to ensure the speculators-format model produces
|
||||
correct outputs, and checks that acceptance length does not collapse under
|
||||
batched inference (lm-eval style).
|
||||
|
||||
Observed per-position acceptance rates on GSM8K (1319 prompts):
|
||||
pos 0: 0.795, pos 1: 0.611, pos 2: 0.429, pos 3: 0.282,
|
||||
pos 4: 0.169, pos 5: 0.093, pos 6: 0.048, pos 7: 0.023
|
||||
Observed mean AL: 3.45 (GSM8K dataset, max_num_seqs=128)
|
||||
"""
|
||||
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
|
||||
spec_llm = LLM(
|
||||
model=MODEL_PATH,
|
||||
model=config.model_path,
|
||||
trust_remote_code=True,
|
||||
max_model_len=4096,
|
||||
max_num_seqs=128,
|
||||
@@ -137,7 +176,7 @@ def test_dflash_speculators_correctness(monkeypatch):
|
||||
|
||||
results = evaluate_gsm8k_offline(spec_llm)
|
||||
accuracy = results["accuracy"]
|
||||
accuracy_threshold = EXPECTED_GSM8K_ACCURACY * (1 - ACCURACY_RTOL)
|
||||
accuracy_threshold = config.expected_gsm8k_accuracy * (1 - config.accuracy_rtol)
|
||||
assert accuracy >= accuracy_threshold, (
|
||||
f"Expected GSM8K accuracy >= {accuracy_threshold:.3f}, got {accuracy:.3f}"
|
||||
)
|
||||
@@ -147,19 +186,18 @@ def test_dflash_speculators_correctness(monkeypatch):
|
||||
print_spec_decode_stats(stats)
|
||||
|
||||
acceptance_len = stats["acceptance_len"]
|
||||
al_threshold = EXPECTED_ACCEPTANCE_LEN * (1 - ACCEPTANCE_LEN_RTOL)
|
||||
al_threshold = config.expected_acceptance_len * (1 - config.acceptance_len_rtol)
|
||||
assert acceptance_len >= al_threshold, (
|
||||
f"DFlash speculators acceptance length too low: "
|
||||
f"{config.display_name} speculators acceptance length too low: "
|
||||
f"{acceptance_len:.2f} < {al_threshold:.2f}"
|
||||
)
|
||||
|
||||
# Check per-position acceptance rates for the first few positions.
|
||||
per_pos_rates = stats["per_pos_acceptance_rates"]
|
||||
for i, expected_rate in enumerate(EXPECTED_PER_POS_ACCEPTANCE_RATES):
|
||||
for i, expected_rate in enumerate(config.expected_per_pos_acceptance_rates):
|
||||
assert i < len(per_pos_rates), (
|
||||
f"Missing per-position acceptance rate for position {i}"
|
||||
)
|
||||
threshold = expected_rate * (1 - PER_POS_RTOL)
|
||||
threshold = expected_rate * (1 - config.per_pos_rtol)
|
||||
assert per_pos_rates[i] >= threshold, (
|
||||
f"Per-position acceptance rate at pos {i} too low: "
|
||||
f"{per_pos_rates[i]:.4f} < {threshold:.4f} "
|
||||
@@ -0,0 +1,85 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Build DeepGEMM's `_C` pybind11 extension for <TARGET_PY>.
|
||||
|
||||
Driven from cmake/external_projects/deepgemm.cmake. The driver runs against
|
||||
the build interpreter's torch; <TARGET_PY> is only consulted for INCLUDEPY
|
||||
and SOABI, so target venvs don't need torch installed.
|
||||
|
||||
Usage: python build_deepgemm_C.py <DEEPGEMM_SRC_DIR> <OUTPUT_DIR> <TARGET_PY>
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from torch.utils import cpp_extension
|
||||
|
||||
if len(sys.argv) != 4:
|
||||
sys.exit(f"usage: {sys.argv[0]} <SRC> <OUT> <TARGET_PY>")
|
||||
|
||||
src = Path(sys.argv[1]).resolve()
|
||||
out = Path(sys.argv[2]).resolve()
|
||||
target_py = sys.argv[3]
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
info = json.loads(
|
||||
subprocess.check_output(
|
||||
[
|
||||
target_py,
|
||||
"-c",
|
||||
"import sysconfig, json; "
|
||||
"print(json.dumps({k: sysconfig.get_config_var(k) "
|
||||
"for k in ('EXT_SUFFIX', 'INCLUDEPY')}))",
|
||||
]
|
||||
).decode()
|
||||
)
|
||||
|
||||
cuda_home = cpp_extension.CUDA_HOME
|
||||
if cuda_home is None:
|
||||
sys.exit("CUDA_HOME not found; cannot build DeepGEMM _C")
|
||||
# CCCL lives outside the standard CUDAToolkit search (mirrors DeepGEMM's setup.py).
|
||||
includes = [
|
||||
info["INCLUDEPY"],
|
||||
f"{cuda_home}/include",
|
||||
f"{cuda_home}/include/cccl",
|
||||
str(src / "csrc"),
|
||||
str(src / "deep_gemm/include"),
|
||||
str(src / "third-party/cutlass/include"),
|
||||
str(src / "third-party/cutlass/tools/util/include"),
|
||||
str(src / "third-party/fmt/include"),
|
||||
*cpp_extension.include_paths(device_type="cuda"),
|
||||
]
|
||||
|
||||
cmd = [
|
||||
os.environ.get("CXX", "g++"),
|
||||
"-shared",
|
||||
"-fPIC",
|
||||
"-std=c++20",
|
||||
"-O3",
|
||||
"-g0",
|
||||
"-Wno-psabi",
|
||||
"-Wno-deprecated-declarations",
|
||||
"-DTORCH_API_INCLUDE_EXTENSION_H",
|
||||
"-DTORCH_EXTENSION_NAME=_C",
|
||||
f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}",
|
||||
*(f"-I{p}" for p in includes),
|
||||
str(src / "csrc/python_api.cpp"),
|
||||
*(f"-L{p}" for p in cpp_extension.library_paths(device_type="cuda")),
|
||||
f"-L{cuda_home}/lib64",
|
||||
"-ltorch",
|
||||
"-ltorch_python",
|
||||
"-ltorch_cpu",
|
||||
"-ltorch_cuda",
|
||||
"-lc10",
|
||||
"-lc10_cuda",
|
||||
"-lcudart",
|
||||
"-lnvrtc",
|
||||
"-o",
|
||||
str(out / f"_C{info['EXT_SUFFIX']}"),
|
||||
]
|
||||
print("[build_deepgemm_C] " + " ".join(cmd), flush=True)
|
||||
subprocess.check_call(cmd)
|
||||
@@ -0,0 +1,41 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Assert the installed vLLM has a `_C.cpython-X.Y-*.so` for every CPython
|
||||
covered by `requires-python`. Fails closed if a Python's `.so` is missing
|
||||
from the wheel — i.e. the regression that surfaced in #41476/#41512.
|
||||
|
||||
Run from a CI test job after vLLM is installed, e.g. the H100 deepgemm
|
||||
kernel tests in .buildkite/test_areas/kernels.yaml.
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import regex as re
|
||||
import tomllib
|
||||
|
||||
SO_RE = re.compile(r"^_C\.cpython-(\d)(\d+)-")
|
||||
|
||||
|
||||
def required_pythons() -> list[str]:
|
||||
pyproject = Path(__file__).resolve().parent.parent / "pyproject.toml"
|
||||
spec = tomllib.loads(pyproject.read_text())["project"]["requires-python"]
|
||||
m = re.match(r">=3\.(\d+),<3\.(\d+)", spec)
|
||||
if not m:
|
||||
sys.exit(f"unexpected requires-python format: {spec!r}")
|
||||
return [f"3.{v}" for v in range(int(m[1]), int(m[2]))]
|
||||
|
||||
|
||||
spec = importlib.util.find_spec("vllm.third_party.deep_gemm")
|
||||
if spec is None or spec.origin is None:
|
||||
sys.exit("vllm.third_party.deep_gemm not importable; is vllm installed?")
|
||||
pkg_dir = Path(spec.origin).parent
|
||||
|
||||
found = {f"{m[1]}.{m[2]}" for f in os.listdir(pkg_dir) if (m := SO_RE.match(f))}
|
||||
required = required_pythons()
|
||||
missing = [v for v in required if v not in found]
|
||||
print(f"deepgemm _C: found {sorted(found)}, required {required}, missing {missing}")
|
||||
sys.exit(1 if missing else 0)
|
||||
Executable
+36
@@ -0,0 +1,36 @@
|
||||
#!/usr/bin/env bash
|
||||
# Provision one bare Python per `requires-python` entry (or per argument) and
|
||||
# print their paths as ":"-separated DEEPGEMM_PYTHON_INTERPRETERS. Skip this
|
||||
# entirely if you already have interpreter paths.
|
||||
#
|
||||
# Usage:
|
||||
# export DEEPGEMM_PYTHON_INTERPRETERS=$(tools/setup_deepgemm_pythons.sh)
|
||||
# python setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
|
||||
#
|
||||
# Optional: DEEPGEMM_VENV_PREFIX (default: /tmp/dgenv).
|
||||
set -euo pipefail
|
||||
|
||||
if [ "$#" -eq 0 ]; then
|
||||
# Derive the matrix from `requires-python = ">=3.X,<3.Y"` in pyproject.toml.
|
||||
pyproject="$(dirname "$0")/../pyproject.toml"
|
||||
spec=$(grep -E '^requires-python' "$pyproject" \
|
||||
| grep -oE '>=3\.[0-9]+,<3\.[0-9]+')
|
||||
lo=${spec#>=3.}; lo=${lo%%,*}
|
||||
hi=${spec##*<3.}
|
||||
set -- $(seq "$lo" $((hi - 1)) | sed 's/^/3./')
|
||||
fi
|
||||
|
||||
prefix="${DEEPGEMM_VENV_PREFIX:-/tmp/dgenv}"
|
||||
mkdir -p "$prefix"
|
||||
|
||||
paths=""
|
||||
for V in "$@"; do
|
||||
venv="$prefix/$V"
|
||||
# uv-managed Python ensures Python.h is present; system 3.X-dev packages
|
||||
# on the manylinux / Ubuntu build bases are not always installed.
|
||||
[ -x "$venv/bin/python" ] || \
|
||||
uv venv --python "$V" "$venv" --python-preference only-managed --seed \
|
||||
>/dev/null
|
||||
paths="$paths:$venv/bin/python"
|
||||
done
|
||||
echo "${paths#:}"
|
||||
@@ -143,6 +143,11 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
|
||||
if self.pass_config.fuse_gemm_comms:
|
||||
self.passes += [AsyncTPPass(config)]
|
||||
|
||||
if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
|
||||
# Run the more specific RMSNorm+router-pad fusion before
|
||||
# AR+RMS, since both consume fused_add_rms_norm.
|
||||
self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
|
||||
|
||||
if self.pass_config.fuse_allreduce_rms:
|
||||
if rocm_aiter_ops.is_enabled():
|
||||
self.passes += [RocmAiterAllReduceFusionPass(config)]
|
||||
@@ -164,9 +169,6 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
|
||||
if rocm_aiter_ops.is_enabled():
|
||||
self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]
|
||||
|
||||
if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
|
||||
self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
|
||||
|
||||
if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled():
|
||||
self.passes += [MLADualRMSNormFusionPass(config)]
|
||||
|
||||
|
||||
@@ -65,8 +65,13 @@ class DeviceConfig:
|
||||
elif isinstance(self.device, torch.device):
|
||||
self.device_type = self.device.type
|
||||
|
||||
# Some device types require processing inputs on CPU
|
||||
if self.device_type in ["tpu"]:
|
||||
# Some platforms require processing inputs on CPU.
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if (
|
||||
current_platform.uses_host_device_handling()
|
||||
and self.device_type == current_platform.device_type
|
||||
):
|
||||
self.device = None
|
||||
else:
|
||||
# Set device with device type
|
||||
|
||||
@@ -1345,7 +1345,7 @@ class ModelConfig:
|
||||
)
|
||||
raise AssertionError(f"Unsupported block type: {block_type}")
|
||||
|
||||
def get_mamba_chunk_size(self) -> int | None:
|
||||
def get_mamba_chunk_size(self) -> int:
|
||||
"""
|
||||
Returns the mamba chunk size if it exists
|
||||
"""
|
||||
@@ -1356,7 +1356,7 @@ class ModelConfig:
|
||||
chunk_size = getattr(self.hf_text_config, "chunk_size", None)
|
||||
|
||||
# Since Mamba1 does not have a chunk notion
|
||||
# we use a default chunk size of 1024.
|
||||
# we use a default chunk size of 2048.
|
||||
if chunk_size is None:
|
||||
chunk_size = 2048
|
||||
|
||||
|
||||
@@ -335,6 +335,7 @@ class NixlEPAll2AllManager(All2AllManagerBase):
|
||||
_lock = threading.Lock()
|
||||
|
||||
def __init__(self, cpu_group, tcp_store_group=None):
|
||||
assert tcp_store_group is not None
|
||||
super().__init__(cpu_group, tcp_store_group)
|
||||
|
||||
self.max_num_ep_ranks = envs.VLLM_NIXL_EP_MAX_NUM_RANKS
|
||||
|
||||
@@ -270,7 +270,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
|
||||
input_tensor = input_.movedim(0, dim).contiguous()
|
||||
|
||||
if sizes is not None:
|
||||
assert len(sizes) == world_size
|
||||
assert len(sizes) == world_size, f"{len(sizes)} == {world_size}"
|
||||
assert input_tensor.shape[0] == sum(sizes)
|
||||
chunk_size = sizes[self.rank_in_group]
|
||||
else:
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user