forked from Karylab-cklius/vllm
Compare commits
152
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c3f8936aca | ||
|
|
ca3e62d336 | ||
|
|
38e16678ba | ||
|
|
27702f6d08 | ||
|
|
22a3cbe152 | ||
|
|
d5b31c954d | ||
|
|
ee38750a75 | ||
|
|
27e0057aed | ||
|
|
f39bcf1e30 | ||
|
|
6467213a9f | ||
|
|
df8e63f4ed | ||
|
|
242afc6bf4 | ||
|
|
5d0fd87038 | ||
|
|
d8deb5b7ad | ||
|
|
2e777d21a8 | ||
|
|
e43a791284 | ||
|
|
66d1cc0c77 | ||
|
|
1c58876618 | ||
|
|
51c1ee9b7c | ||
|
|
213f10bfdd | ||
|
|
e87e09a50a | ||
|
|
809b98e5b7 | ||
|
|
b53c507bc9 | ||
|
|
2d7d6cf765 | ||
|
|
91740ca5ea | ||
|
|
e47c98ef7a | ||
|
|
aee190ac37 | ||
|
|
16e336491e | ||
|
|
c7aa186d67 | ||
|
|
f653761252 | ||
|
|
4a8ae26e53 | ||
|
|
1333864408 | ||
|
|
01b9b5af67 | ||
|
|
8c57b6e7bc | ||
|
|
79246b5ea6 | ||
|
|
48954de237 | ||
|
|
c6235ed180 | ||
|
|
628c436301 | ||
|
|
2228fe6868 | ||
|
|
84bd8a3c1e | ||
|
|
b786ec8e74 | ||
|
|
20dcd984f9 | ||
|
|
6fca518157 | ||
|
|
98661fe012 | ||
|
|
b0765bee17 | ||
|
|
0a201b60cf | ||
|
|
8b9ea2f881 | ||
|
|
2ceea42958 | ||
|
|
bee126165f | ||
|
|
27cc676be3 | ||
|
|
4845aee6b7 | ||
|
|
0c620d2e08 | ||
|
|
6bb924bbf3 | ||
|
|
eaec7be446 | ||
|
|
f04fd1677b | ||
|
|
420b0a5c95 | ||
|
|
1e9500410a | ||
|
|
416f9cdede | ||
|
|
685bf811d6 | ||
|
|
e1e4646b06 | ||
|
|
4f2af1a7c0 | ||
|
|
577b9623e6 | ||
|
|
1cb0838721 | ||
|
|
be5983b874 | ||
|
|
9c07342fdc | ||
|
|
844df54269 | ||
|
|
422dd02598 | ||
|
|
8c780943b4 | ||
|
|
e724b0ea8d | ||
|
|
712ad0286c | ||
|
|
321fa2d6d1 | ||
|
|
3e1ad4435f | ||
|
|
8decbfa02c | ||
|
|
62ba7516e8 | ||
|
|
6f53753fc9 | ||
|
|
6ec9bbec38 | ||
|
|
01d4d1ad37 | ||
|
|
c103c02a1a | ||
|
|
67058ca326 | ||
|
|
894a02500b | ||
|
|
66dfee7121 | ||
|
|
db9a84e0cd | ||
|
|
cb03fee32b | ||
|
|
c51df43005 | ||
|
|
54dc64d5d3 | ||
|
|
e6ff3e9c83 | ||
|
|
08834cc3ce | ||
|
|
856ec4804a | ||
|
|
1c607d7b2c | ||
|
|
4f7309fcc0 | ||
|
|
0a9362d6ab | ||
|
|
cfd2573f23 | ||
|
|
c3ad791e1a | ||
|
|
8586369f61 | ||
|
|
ae3b4deb8a | ||
|
|
c293ccc58e | ||
|
|
d58c42e19c | ||
|
|
3e49479c4b | ||
|
|
964a4bc2a5 | ||
|
|
c408fdd663 | ||
|
|
5737770c6c | ||
|
|
0c99629ede | ||
|
|
edd60ac93a | ||
|
|
bcf5cac9fb | ||
|
|
a9484dac7b | ||
|
|
f3fef12350 | ||
|
|
51295793a2 | ||
|
|
3ccc1ff495 | ||
|
|
aeb92c37ab | ||
|
|
529c671e80 | ||
|
|
bc635fad23 | ||
|
|
c3e64696cd | ||
|
|
4f7bde572a | ||
|
|
2fa1f8ec00 | ||
|
|
7075df79b3 | ||
|
|
0dbaf9daad | ||
|
|
a3ec4a35f5 | ||
|
|
32964e7700 | ||
|
|
a07642667d | ||
|
|
c3868bbbe4 | ||
|
|
947138b6c2 | ||
|
|
941fb50835 | ||
|
|
6b6ac6c3c7 | ||
|
|
b542bdf7fb | ||
|
|
415a879899 | ||
|
|
7198940b39 | ||
|
|
14043dfecd | ||
|
|
1adaa5056b | ||
|
|
4d5c89295b | ||
|
|
dd5506a157 | ||
|
|
a3c83ff2fd | ||
|
|
9c61864bf8 | ||
|
|
71725f6730 | ||
|
|
b4806c8ee1 | ||
|
|
526927be94 | ||
|
|
75a4c166f2 | ||
|
|
2917d6363a | ||
|
|
bf3a797a1d | ||
|
|
6ba036fb9b | ||
|
|
e23e0311a3 | ||
|
|
7e8c3db5fb | ||
|
|
f6a456de8c | ||
|
|
14f74bd812 | ||
|
|
275cf255f7 | ||
|
|
18a428b606 | ||
|
|
239e6ff95b | ||
|
|
19844eecbf | ||
|
|
057cc9c0e5 | ||
|
|
1dfb431175 | ||
|
|
76de80cebd | ||
|
|
798de5d4be | ||
|
|
5a21f7d139 |
@@ -14,13 +14,15 @@ steps:
|
||||
- tests/kernels/moe/test_cpu_fused_moe.py
|
||||
- tests/kernels/test_onednn.py
|
||||
- tests/kernels/test_awq_int4_to_int8.py
|
||||
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 20m "
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
|
||||
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
|
||||
pytest -x -v -s tests/kernels/test_onednn.py
|
||||
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py"
|
||||
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
|
||||
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py"
|
||||
|
||||
- label: CPU-Compatibility Tests
|
||||
depends_on: []
|
||||
|
||||
@@ -37,7 +37,7 @@ steps:
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
@@ -76,7 +76,7 @@ steps:
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
@@ -309,6 +309,7 @@ steps:
|
||||
depends_on: ~
|
||||
|
||||
- label: "Build release image - x86_64 - CPU"
|
||||
key: build-cpu-release-image-x86
|
||||
depends_on:
|
||||
- block-cpu-release-image-build
|
||||
- input-release-version
|
||||
@@ -327,7 +328,8 @@ steps:
|
||||
depends_on: ~
|
||||
|
||||
- label: "Build release image - arm64 - CPU"
|
||||
depends_on:
|
||||
key: build-cpu-release-image-arm64
|
||||
depends_on:
|
||||
- block-arm64-cpu-release-image-build
|
||||
- input-release-version
|
||||
agents:
|
||||
@@ -436,6 +438,41 @@ 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:
|
||||
@@ -723,7 +760,7 @@ steps:
|
||||
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
|
||||
env:
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
VARIANT: "rocm721"
|
||||
VARIANT: "rocm722"
|
||||
|
||||
# ROCm Job 6: Build ROCm Release Docker Image
|
||||
- label: ":docker: Build release image - x86_64 - ROCm"
|
||||
|
||||
@@ -8,8 +8,6 @@ if [ -z "${RELEASE_VERSION}" ]; then
|
||||
RELEASE_VERSION="1.0.0.dev"
|
||||
fi
|
||||
|
||||
ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
|
||||
|
||||
buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
|
||||
To download the wheel (by commit):
|
||||
\`\`\`
|
||||
@@ -25,95 +23,5 @@ aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
\`\`\`
|
||||
|
||||
|
||||
To download and upload the image:
|
||||
|
||||
\`\`\`
|
||||
# Download images:
|
||||
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
|
||||
|
||||
# Tag and push images:
|
||||
|
||||
## CUDA
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
|
||||
docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64
|
||||
docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
docker push vllm/vllm-openai:latest-x86_64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 vllm/vllm-openai:x86_64-cu129
|
||||
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
|
||||
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
|
||||
docker push vllm/vllm-openai:latest-x86_64-cu129
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
|
||||
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
|
||||
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker push vllm/vllm-openai:latest-aarch64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 vllm/vllm-openai:aarch64-cu129
|
||||
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
docker push vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
|
||||
## ROCm
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
docker push vllm/vllm-openai-rocm:latest
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
|
||||
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
docker push vllm/vllm-openai-rocm:latest-base
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
|
||||
## CPU
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:x86_64
|
||||
docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:latest-x86_64
|
||||
docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
|
||||
docker push vllm/vllm-openai-cpu:latest-x86_64
|
||||
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:arm64
|
||||
docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:latest-arm64
|
||||
docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
|
||||
docker push vllm/vllm-openai-cpu:latest-arm64
|
||||
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
|
||||
|
||||
# Create multi-arch manifest:
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest
|
||||
docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker manifest push vllm/vllm-openai:latest
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest-cu129
|
||||
docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
docker manifest push vllm/vllm-openai:latest-cu129
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
|
||||
|
||||
docker manifest rm vllm/vllm-openai-cpu:latest || true
|
||||
docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
|
||||
docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
|
||||
docker manifest push vllm/vllm-openai-cpu:latest
|
||||
docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
|
||||
\`\`\`
|
||||
Docker images are published automatically by the "Publish release images to DockerHub" pipeline step.
|
||||
EOF
|
||||
|
||||
Executable
+55
@@ -0,0 +1,55 @@
|
||||
#!/bin/bash
|
||||
# Usage: ./ci-fetch-log.sh <buildkite_job_url> [output_file]
|
||||
# ./ci-fetch-log.sh <build_number> <job_uuid> [output_file]
|
||||
#
|
||||
# Downloads the raw log for a Buildkite job from the public, unauthenticated
|
||||
# /organizations/<org>/pipelines/<pipeline>/builds/<n>/jobs/<uuid>/download
|
||||
# endpoint, then strips ANSI/timestamps via ci-clean-log.sh.
|
||||
#
|
||||
# Find <build_number> and <job_uuid> via:
|
||||
# gh pr checks <PR> --repo vllm-project/vllm
|
||||
# Each failing row's URL is .../builds/<build_number>#<job_uuid>.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
ORG="vllm"
|
||||
PIPELINE="ci"
|
||||
|
||||
usage() {
|
||||
echo "Usage: $0 <buildkite_job_url> [output_file]"
|
||||
echo " $0 <build_number> <job_uuid> [output_file]"
|
||||
exit 1
|
||||
}
|
||||
|
||||
if [ $# -lt 1 ]; then usage; fi
|
||||
|
||||
if [[ "$1" == https://* ]]; then
|
||||
BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p')
|
||||
JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1)
|
||||
OUT="${2:-ci-${BUILD}-${JOB:0:8}.log}"
|
||||
else
|
||||
if [ $# -lt 2 ]; then usage; fi
|
||||
BUILD="$1"
|
||||
JOB="$2"
|
||||
OUT="${3:-ci-${BUILD}-${JOB:0:8}.log}"
|
||||
fi
|
||||
|
||||
if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
|
||||
echo "Could not parse build number or job UUID from: $1" >&2
|
||||
usage
|
||||
fi
|
||||
|
||||
COOKIES=$(mktemp)
|
||||
trap 'rm -f "$COOKIES"' EXIT
|
||||
|
||||
# Buildkite issues a session cookie on first hit; subsequent /download needs it.
|
||||
curl -fsSL -c "$COOKIES" -A "vllm-ci-fetch-log" \
|
||||
"https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}" -o /dev/null
|
||||
|
||||
curl -fsSL -b "$COOKIES" -A "vllm-ci-fetch-log" \
|
||||
"https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/${JOB}/download" \
|
||||
-o "$OUT"
|
||||
|
||||
bash "$(dirname "$0")/ci-clean-log.sh" "$OUT"
|
||||
|
||||
echo "$OUT"
|
||||
@@ -378,9 +378,11 @@ HF_MOUNT="/root/.cache/huggingface"
|
||||
# double-quotes will have been stripped by the calling shell.
|
||||
if [[ -n "${VLLM_TEST_COMMANDS:-}" ]]; then
|
||||
commands="${VLLM_TEST_COMMANDS}"
|
||||
commands_source="env"
|
||||
echo "Commands sourced from VLLM_TEST_COMMANDS (quoting preserved)"
|
||||
else
|
||||
commands="$*"
|
||||
commands_source="argv"
|
||||
if [[ -z "$commands" ]]; then
|
||||
echo "Error: No test commands provided." >&2
|
||||
echo "Usage:" >&2
|
||||
@@ -397,9 +399,15 @@ fi
|
||||
|
||||
echo "Raw commands: $commands"
|
||||
|
||||
# Fix quoting before ROCm overrides (so overrides see correct structure)
|
||||
commands=$(re_quote_pytest_markers "$commands")
|
||||
echo "After re-quoting: $commands"
|
||||
# Only try to repair stripped pytest -m/-k quoting in legacy argv mode.
|
||||
# VLLM_TEST_COMMANDS preserves inner quoting already, and re-quoting that path
|
||||
# can corrupt embedded echo strings or otherwise well-formed shell fragments.
|
||||
if [[ "$commands_source" == "argv" ]]; then
|
||||
commands=$(re_quote_pytest_markers "$commands")
|
||||
echo "After re-quoting: $commands"
|
||||
else
|
||||
echo "Skipping re-quoting for VLLM_TEST_COMMANDS input"
|
||||
fi
|
||||
|
||||
commands=$(apply_rocm_test_overrides "$commands")
|
||||
echo "Final commands: $commands"
|
||||
|
||||
Executable
+180
@@ -0,0 +1,180 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Publish release Docker images from ECR to DockerHub.
|
||||
# Pulls per-arch images, tags with latest and versioned tags, pushes them,
|
||||
# then creates and pushes multi-arch manifests.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
|
||||
if [ -z "${RELEASE_VERSION}" ]; then
|
||||
echo "ERROR: release-version metadata not set"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
COMMIT="$BUILDKITE_COMMIT"
|
||||
ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
|
||||
|
||||
echo "========================================"
|
||||
echo "Publishing release images v${RELEASE_VERSION}"
|
||||
echo " Commit: ${COMMIT}"
|
||||
echo " ROCm base cache key: ${ROCM_BASE_CACHE_KEY}"
|
||||
echo "========================================"
|
||||
|
||||
# Login to ECR to pull staging images
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
# ---- CUDA (default: 13.0) ----
|
||||
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
docker push vllm/vllm-openai:latest-x86_64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker push vllm/vllm-openai:latest-aarch64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest || true
|
||||
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true
|
||||
docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker manifest push vllm/vllm-openai:latest
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
|
||||
|
||||
# ---- CUDA 12.9 ----
|
||||
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
|
||||
docker push vllm/vllm-openai:latest-x86_64-cu129
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
docker push vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest-cu129 || true
|
||||
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true
|
||||
docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
docker manifest push vllm/vllm-openai:latest-cu129
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
|
||||
|
||||
# ---- Ubuntu 24.04 (CUDA 13.0) ----
|
||||
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404
|
||||
docker push vllm/vllm-openai:latest-x86_64-ubuntu2404
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
|
||||
docker push vllm/vllm-openai:latest-aarch64-ubuntu2404
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true
|
||||
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true
|
||||
docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
|
||||
docker manifest push vllm/vllm-openai:latest-ubuntu2404
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404
|
||||
|
||||
# ---- Ubuntu 24.04 (CUDA 12.9) ----
|
||||
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404
|
||||
docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
|
||||
docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true
|
||||
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true
|
||||
docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
|
||||
docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404
|
||||
|
||||
# ---- ROCm ----
|
||||
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
docker push vllm/vllm-openai-rocm:latest
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
docker push vllm/vllm-openai-rocm:latest-base
|
||||
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
|
||||
|
||||
# ---- CPU ----
|
||||
# CPU images are behind separate block steps and may not have been built.
|
||||
# All-or-nothing: inspect both arches first, then either publish everything
|
||||
# (per-arch + multi-arch manifest) or skip everything. Publishing only one
|
||||
# arch would leave `:latest-x86_64` pointing at the new release while the
|
||||
# `:latest` multi-arch manifest still resolves to the previous release.
|
||||
|
||||
CPU_X86_TAG=public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
|
||||
CPU_ARM_TAG=public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
|
||||
|
||||
CPU_X86_AVAILABLE=false
|
||||
CPU_ARM_AVAILABLE=false
|
||||
docker manifest inspect "${CPU_X86_TAG}" >/dev/null 2>&1 && CPU_X86_AVAILABLE=true
|
||||
docker manifest inspect "${CPU_ARM_TAG}" >/dev/null 2>&1 && CPU_ARM_AVAILABLE=true
|
||||
|
||||
if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then
|
||||
docker pull "${CPU_X86_TAG}"
|
||||
docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64
|
||||
docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
|
||||
docker push vllm/vllm-openai-cpu:latest-x86_64
|
||||
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
|
||||
|
||||
docker pull "${CPU_ARM_TAG}"
|
||||
docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64
|
||||
docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
|
||||
docker push vllm/vllm-openai-cpu:latest-arm64
|
||||
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
|
||||
|
||||
docker manifest rm vllm/vllm-openai-cpu:latest || true
|
||||
docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true
|
||||
docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
|
||||
docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
|
||||
docker manifest push vllm/vllm-openai-cpu:latest
|
||||
docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
|
||||
elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then
|
||||
echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)"
|
||||
else
|
||||
# Partial state: one arch built, the other did not. Fail loudly rather than
|
||||
# ship a Docker Hub state where `:latest-${arch}` and `:latest` (multi-arch)
|
||||
# disagree on which release they point at.
|
||||
echo "ERROR: Partial CPU build detected (x86_64=${CPU_X86_AVAILABLE}, arm64=${CPU_ARM_AVAILABLE})."
|
||||
echo " Refusing to publish to avoid split-tag drift between per-arch and multi-arch tags."
|
||||
echo " Re-run the missing CPU build and retry, or manually publish if a single-arch release is intended."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "Successfully published release images for v${RELEASE_VERSION}"
|
||||
@@ -51,6 +51,7 @@ vllm serve "$MODEL" \
|
||||
--offload-num-in-group 2 \
|
||||
--offload-prefetch-step 1 \
|
||||
--offload-params w13_weight w2_weight \
|
||||
--generation-config vllm \
|
||||
--port "$PORT" \
|
||||
${EXTRA_ARGS+"${EXTRA_ARGS[@]}"} &
|
||||
SERVER_PID=$!
|
||||
|
||||
@@ -39,10 +39,11 @@ fi
|
||||
|
||||
set -x # avoid printing secrets above
|
||||
|
||||
# install twine from pypi
|
||||
# install twine and sdist build prerequisites from pypi
|
||||
python3 -m venv /tmp/vllm-release-env
|
||||
source /tmp/vllm-release-env/bin/activate
|
||||
pip install twine
|
||||
pip install -r requirements/build/cuda.txt
|
||||
python3 -m twine --version
|
||||
|
||||
# copy release wheels to local directory
|
||||
|
||||
@@ -621,6 +621,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -1108,6 +1109,7 @@ steps:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
|
||||
- pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
|
||||
- pytest -v -s tests/compile/passes/distributed/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions
|
||||
|
||||
#----------------------------------------------------------- mi300 · cuda ------------------------------------------------------------#
|
||||
|
||||
@@ -1802,6 +1804,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -1843,6 +1846,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2203,7 +2207,6 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2627,6 +2630,7 @@ steps:
|
||||
agent_pool: mi325_1
|
||||
torch_nightly: true
|
||||
parallelism: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2652,6 +2656,7 @@ steps:
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi325]
|
||||
agent_pool: mi325_1
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2748,6 +2753,7 @@ steps:
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -2763,6 +2769,7 @@ steps:
|
||||
agent_pool: mi355_1
|
||||
fast_check: true
|
||||
torch_nightly: true
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -3059,6 +3066,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -3258,6 +3266,7 @@ steps:
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -3283,6 +3292,7 @@ steps:
|
||||
timeout_in_minutes: 60
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
|
||||
agent_pool: mi355_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -17,7 +17,7 @@ steps:
|
||||
- label: V1 attention (B200)
|
||||
key: v1-attention-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
source_file_dependencies:
|
||||
- vllm/config/attention.py
|
||||
- vllm/model_executor/layers/attention
|
||||
|
||||
@@ -14,7 +14,7 @@ steps:
|
||||
|
||||
- label: Attention Benchmarks Smoke Test (B200)
|
||||
key: attention-benchmarks-smoke-test-b200
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/"
|
||||
|
||||
@@ -43,7 +43,7 @@ steps:
|
||||
key: asynctp-correctness-tests-b200
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
@@ -68,7 +68,7 @@ steps:
|
||||
key: fusion-and-compile-unit-tests-2xb200
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/fp4/
|
||||
- vllm/model_executor/layers/quantization/
|
||||
@@ -137,7 +137,7 @@ steps:
|
||||
key: fusion-e2e-config-sweep-b200
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
num_devices: 1
|
||||
optional: true
|
||||
commands:
|
||||
@@ -209,7 +209,7 @@ steps:
|
||||
key: fusion-e2e-tp2-b200
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/
|
||||
|
||||
@@ -8,7 +8,7 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
@@ -19,7 +19,7 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
@@ -31,7 +31,7 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
@@ -43,7 +43,7 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
@@ -55,7 +55,7 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
@@ -67,7 +67,7 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
@@ -83,7 +83,7 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- vllm/v1/worker/kv_connector_model_runner_mixin.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
@@ -96,7 +96,7 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
|
||||
@@ -212,7 +212,7 @@ steps:
|
||||
|
||||
- label: Distributed Tests (2 GPUs)(B200)
|
||||
key: distributed-tests-2-gpus-b200
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/"
|
||||
num_devices: 2
|
||||
|
||||
@@ -25,7 +25,7 @@ steps:
|
||||
- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
|
||||
key: qwen3-30b-a3b-fp8-block-accuracy-b200
|
||||
timeout_in_minutes: 60
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 2
|
||||
working_dir: "/vllm-workspace"
|
||||
|
||||
@@ -13,8 +13,9 @@ steps:
|
||||
- tests/test_config
|
||||
- tests/test_logger
|
||||
- tests/test_vllm_port
|
||||
- tests/test_jit_monitor.py
|
||||
commands:
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
|
||||
|
||||
- label: Engine (1 GPU)
|
||||
key: engine-1-gpu
|
||||
@@ -56,11 +57,6 @@ steps:
|
||||
commands:
|
||||
# Only run tests that need exactly 2 GPUs
|
||||
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "tensor_parallelism"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_2
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 e2e (4 GPUs)
|
||||
key: v1-e2e-4-gpus
|
||||
@@ -73,11 +69,6 @@ steps:
|
||||
commands:
|
||||
# Only run tests that need 4 GPUs
|
||||
- pytest -v -s v1/e2e/spec_decode/test_spec_decode.py -k "eagle_correctness_heavy"
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_4
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 e2e (4xH100)
|
||||
key: v1-e2e-4xh100
|
||||
|
||||
@@ -26,6 +26,11 @@ steps:
|
||||
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py
|
||||
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
|
||||
- pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
|
||||
mirror:
|
||||
amd:
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1)
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
@@ -38,11 +43,6 @@ steps:
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s entrypoints/openai/chat_completion --ignore=entrypoints/openai/chat_completion/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/chat_completion/test_oot_registration.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2)
|
||||
@@ -57,11 +57,6 @@ steps:
|
||||
- pytest -v -s entrypoints/openai/completion --ignore=entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- pytest -v -s entrypoints/openai/speech_to_text/
|
||||
- pytest -v -s entrypoints/test_chat_utils.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3)
|
||||
key: entrypoints-integration-api-server-openai-part-3
|
||||
|
||||
@@ -125,7 +125,7 @@ steps:
|
||||
key: kernels-b200
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
# optional: true
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/fp4/
|
||||
@@ -212,7 +212,7 @@ steps:
|
||||
- label: Kernels Fp4 MoE Test (B200)
|
||||
key: kernels-fp4-moe-test-b200
|
||||
timeout_in_minutes: 60
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
num_devices: 1
|
||||
optional: true
|
||||
commands:
|
||||
|
||||
@@ -51,7 +51,7 @@ steps:
|
||||
- label: LM Eval Qwen3.5 Models (B200)
|
||||
key: lm-eval-qwen3-5-models-b200
|
||||
timeout_in_minutes: 120
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
@@ -84,7 +84,7 @@ steps:
|
||||
|
||||
- label: MoE Refactor Integration Test (B200 - TEMPORARY)
|
||||
key: moe-refactor-integration-test-b200-temporary
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 2
|
||||
commands:
|
||||
@@ -137,3 +137,10 @@ steps:
|
||||
commands:
|
||||
- uv pip install --system 'gpt-oss[eval]==0.0.5'
|
||||
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
|
||||
|
||||
- label: MRCR Eval Small Models
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
- tests/evals/mrcr/
|
||||
commands:
|
||||
- pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt
|
||||
|
||||
@@ -12,11 +12,6 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# TODO: create another `optional` test group for slow tests
|
||||
- pytest -v -s -m 'not slow_test' v1/spec_decode
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 Sample + Logits
|
||||
key: v1-sample-logits
|
||||
@@ -38,7 +33,7 @@ steps:
|
||||
- pytest -v -s v1/test_outputs.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -67,11 +62,6 @@ steps:
|
||||
# Integration test for streaming correctness (requires special branch).
|
||||
- pip install -U git+https://github.com/robertgshaw2-redhat/lm-evaluation-harness.git@streaming-api
|
||||
- pytest -v -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 Others (CPU)
|
||||
key: v1-others-cpu
|
||||
@@ -224,7 +214,7 @@ steps:
|
||||
- label: Batch Invariance (B200)
|
||||
key: batch-invariance-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
source_file_dependencies:
|
||||
- vllm/v1/attention
|
||||
- vllm/model_executor/layers
|
||||
|
||||
@@ -42,12 +42,6 @@ steps:
|
||||
- tests/models/test_registry.py
|
||||
commands:
|
||||
- pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
- label: Basic Models Test (Other CPU) # 5min
|
||||
key: basic-models-test-other-cpu
|
||||
|
||||
@@ -48,6 +48,14 @@ steps:
|
||||
parallelism: 2
|
||||
mirror:
|
||||
torch_nightly: {}
|
||||
amd:
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
|
||||
- label: Language Models Test (Extended Generation) # 80min
|
||||
key: language-models-test-extended-generation
|
||||
@@ -62,15 +70,6 @@ steps:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/AndreasKaratzas/mamba@fix-rocm-7.0-warp-size-constexpr'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
|
||||
|
||||
- label: Language Models Test (PPL)
|
||||
key: language-models-test-ppl
|
||||
@@ -92,11 +91,6 @@ steps:
|
||||
- tests/models/language/pooling
|
||||
commands:
|
||||
- pytest -v -s models/language/pooling -m 'not core_model'
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Language Models Test (MTEB)
|
||||
key: language-models-test-mteb
|
||||
|
||||
@@ -15,7 +15,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_ultravox.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -33,7 +33,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -49,7 +49,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_vl.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
@@ -64,11 +64,6 @@ steps:
|
||||
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/processing
|
||||
- pytest models/multimodal/generation/test_memory_leak.py -m core_model
|
||||
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Multi-Modal Processor (CPU)
|
||||
key: multi-modal-processor-cpu
|
||||
@@ -120,7 +115,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/test_mapping.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi300_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
|
||||
@@ -25,7 +25,7 @@ steps:
|
||||
key: quantized-moe-test-b200
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
source_file_dependencies:
|
||||
- tests/quantization/test_blackwell_moe.py
|
||||
- vllm/model_executor/models/deepseek_v2.py
|
||||
|
||||
@@ -17,7 +17,7 @@ steps:
|
||||
- VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
device: mi250_1
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
|
||||
@@ -75,7 +75,7 @@ steps:
|
||||
- label: Spec Decode Draft Model Nightly B200
|
||||
key: spec-decode-draft-model-nightly-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
|
||||
+1
-3
@@ -477,9 +477,7 @@ pull_request_rules:
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^examples/online_serving/disaggregated[^/]*/.*
|
||||
- files~=^examples/offline_inference/disaggregated[^/]*/.*
|
||||
- files~=^examples/others/lmcache/
|
||||
- files~=^examples/disaggregated/
|
||||
- files~=^tests/v1/kv_connector/
|
||||
- files~=^vllm/distributed/kv_transfer/
|
||||
- title~=(?i)\bP/?D\b
|
||||
|
||||
+6
-7
@@ -307,12 +307,12 @@ set(VLLM_EXT_SRC
|
||||
"csrc/quantization/activation_kernels.cu"
|
||||
"csrc/cuda_utils_kernels.cu"
|
||||
"csrc/custom_all_reduce.cu"
|
||||
"csrc/torch_bindings.cpp")
|
||||
"csrc/torch_bindings.cpp"
|
||||
"csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_EXT_SRC
|
||||
"csrc/minimax_reduce_rms_kernel.cu"
|
||||
"csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
|
||||
"csrc/minimax_reduce_rms_kernel.cu")
|
||||
|
||||
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
|
||||
|
||||
@@ -1047,14 +1047,13 @@ endif()
|
||||
set(VLLM_MOE_EXT_SRC
|
||||
"csrc/moe/torch_bindings.cpp"
|
||||
"csrc/moe/moe_align_sum_kernels.cu"
|
||||
"csrc/moe/topk_softmax_kernels.cu")
|
||||
"csrc/moe/topk_softmax_kernels.cu"
|
||||
"csrc/moe/topk_softplus_sqrt_kernels.cu")
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_MOE_EXT_SRC
|
||||
"csrc/moe/moe_wna16.cu"
|
||||
"csrc/moe/grouped_topk_kernels.cu"
|
||||
"csrc/moe/router_gemm.cu"
|
||||
"csrc/moe/topk_softplus_sqrt_kernels.cu")
|
||||
"csrc/moe/grouped_topk_kernels.cu")
|
||||
endif()
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
|
||||
@@ -1473,6 +1473,12 @@ async def main() -> None:
|
||||
"(for example: --warmup-percentages=0%%,50%%)",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--trust-remote-code",
|
||||
action="store_true",
|
||||
help="Trust remote code when loading the tokenizer.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
logger.info(args)
|
||||
@@ -1515,7 +1521,9 @@ async def main() -> None:
|
||||
np.random.seed(args.seed)
|
||||
|
||||
logger.info("Loading tokenizer")
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model)
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
args.model, trust_remote_code=args.trust_remote_code
|
||||
)
|
||||
|
||||
await get_server_info(args.url)
|
||||
|
||||
|
||||
+24
-20
@@ -32,18 +32,23 @@ else()
|
||||
"-DVLLM_CPU_EXTENSION")
|
||||
|
||||
# locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
|
||||
# and create a local shim dir with it
|
||||
# and create a local shim dir with it. When PyTorch is built from source or packaged
|
||||
# by a distro (common on RISC-V, s390x, Fedora/RHEL aarch64), no vendored libgomp
|
||||
# exists and the shim dir is empty; fall back to the system libgomp in that case.
|
||||
vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
|
||||
|
||||
find_library(OPEN_MP
|
||||
NAMES gomp
|
||||
PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
|
||||
NO_DEFAULT_PATH
|
||||
REQUIRED
|
||||
)
|
||||
# Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
|
||||
if (OPEN_MP)
|
||||
if(VLLM_TORCH_GOMP_SHIM_DIR)
|
||||
find_library(OPEN_MP
|
||||
NAMES gomp
|
||||
PATHS "${VLLM_TORCH_GOMP_SHIM_DIR}"
|
||||
NO_DEFAULT_PATH
|
||||
REQUIRED
|
||||
)
|
||||
# Use the same libgomp as PyTorch at runtime
|
||||
set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
|
||||
else()
|
||||
# Fall back to system / toolchain libgomp
|
||||
find_library(OPEN_MP NAMES gomp REQUIRED)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
@@ -321,14 +326,6 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
|
||||
set(ONEDNN_VERBOSE "ON")
|
||||
set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
|
||||
|
||||
# TODO: Refactor this
|
||||
if (ENABLE_X86_ISA)
|
||||
# Note: only enable oneDNN for AVX512
|
||||
list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS_AVX512})
|
||||
else()
|
||||
list(APPEND DNNL_COMPILE_FLAGS ${CXX_COMPILE_FLAGS})
|
||||
endif()
|
||||
|
||||
set(VLLM_BUILD_TYPE ${CMAKE_BUILD_TYPE})
|
||||
set(CMAKE_BUILD_TYPE "Release") # remove oneDNN debug symbols to reduce size
|
||||
FetchContent_MakeAvailable(oneDNN)
|
||||
@@ -341,8 +338,14 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
|
||||
PRIVATE ${oneDNN_SOURCE_DIR}/src
|
||||
)
|
||||
target_link_libraries(dnnl_ext dnnl torch)
|
||||
target_compile_options(dnnl_ext PRIVATE ${DNNL_COMPILE_FLAGS} -fPIC)
|
||||
if (ENABLE_X86_ISA)
|
||||
target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS_AVX2} -fPIC)
|
||||
else()
|
||||
target_compile_options(dnnl_ext PRIVATE ${CXX_COMPILE_FLAGS} -fPIC)
|
||||
endif()
|
||||
list(APPEND LIBS dnnl_ext)
|
||||
|
||||
|
||||
set(USE_ONEDNN ON)
|
||||
else()
|
||||
set(USE_ONEDNN OFF)
|
||||
@@ -430,10 +433,11 @@ if (ENABLE_X86_ISA)
|
||||
"csrc/cpu/pos_encoding.cpp"
|
||||
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp")
|
||||
|
||||
set(VLLM_EXT_SRC_AVX2
|
||||
set(VLLM_EXT_SRC_AVX2
|
||||
"csrc/cpu/utils.cpp"
|
||||
"csrc/cpu/spec_decode_utils.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/dnnl_kernels.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp"
|
||||
# TODO: Remove these files
|
||||
"csrc/cpu/activation.cpp"
|
||||
@@ -448,7 +452,7 @@ if (ENABLE_X86_ISA)
|
||||
|
||||
set(_C_LIBS numa dnnl_ext)
|
||||
set(_C_AVX512_LIBS numa dnnl_ext)
|
||||
set(_C_AVX2_LIBS numa)
|
||||
set(_C_AVX2_LIBS numa dnnl_ext)
|
||||
|
||||
# AMX + AVX512F + AVX512BF16 + AVX512VNNI
|
||||
define_extension_target(
|
||||
|
||||
@@ -29,6 +29,8 @@ torch::Tensor get_scheduler_metadata(
|
||||
isa = cpu_attention::ISA::NEON;
|
||||
} else if (isa_hint == "vxe") {
|
||||
isa = cpu_attention::ISA::VXE;
|
||||
} else if (isa_hint == "vsx") {
|
||||
isa = cpu_attention::ISA::VSX;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
|
||||
}
|
||||
@@ -129,6 +131,8 @@ void cpu_attn_reshape_and_cache(
|
||||
return cpu_attention::ISA::NEON;
|
||||
} else if (isa == "vxe") {
|
||||
return cpu_attention::ISA::VXE;
|
||||
} else if (isa == "vsx") {
|
||||
return cpu_attention::ISA::VSX;
|
||||
} else {
|
||||
TORCH_CHECK(false, "Invalid ISA type: " + isa);
|
||||
}
|
||||
|
||||
@@ -12,7 +12,7 @@
|
||||
#include "cpu/utils.hpp"
|
||||
|
||||
namespace cpu_attention {
|
||||
enum class ISA { AMX, VEC, VEC16, NEON, VXE };
|
||||
enum class ISA { AMX, VEC, VEC16, NEON, VXE, VSX };
|
||||
|
||||
// Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
|
||||
enum class Fp8KVCacheDataType {
|
||||
@@ -164,6 +164,9 @@ struct AttentionMetadata {
|
||||
case ISA::VXE:
|
||||
ss << "VXE, ";
|
||||
break;
|
||||
case ISA::VSX:
|
||||
ss << "VSX, ";
|
||||
break;
|
||||
}
|
||||
ss << "workitem_group_num: " << workitem_group_num
|
||||
<< ", reduction_item_num: " << reduction_item_num
|
||||
|
||||
@@ -27,8 +27,8 @@ FORCE_INLINE std::pair<vec_op::FP32Vec16, vec_op::FP32Vec16> load_b_pair_vec(
|
||||
return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
|
||||
} else {
|
||||
using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
|
||||
return {vec_op::FP32Vec16(load_vec_t(ptr)),
|
||||
vec_op::FP32Vec16(load_vec_t(ptr + 16))};
|
||||
return std::make_pair(vec_op::FP32Vec16(load_vec_t(ptr)),
|
||||
vec_op::FP32Vec16(load_vec_t(ptr + 16)));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,359 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#ifndef CPU_ATTN_VSX_HPP
|
||||
#define CPU_ATTN_VSX_HPP
|
||||
|
||||
#include "cpu_attn_impl.hpp"
|
||||
#include <altivec.h>
|
||||
#include <type_traits>
|
||||
|
||||
namespace cpu_attention {
|
||||
|
||||
namespace {
|
||||
|
||||
// ppc64le Vector = 16 bytes (128 bits)
|
||||
#define BLOCK_SIZE_ALIGNMENT 32
|
||||
#define HEAD_SIZE_ALIGNMENT 32
|
||||
#define MAX_Q_HEAD_NUM_PER_ITER 16
|
||||
|
||||
template <typename kv_cache_t>
|
||||
FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0,
|
||||
__vector float& b1);
|
||||
|
||||
// [1] Float Specialization
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, __vector float& b0,
|
||||
__vector float& b1) {
|
||||
b0 = vec_xl(0, const_cast<float*>(p));
|
||||
b1 = vec_xl(0, const_cast<float*>(p + 4));
|
||||
}
|
||||
|
||||
// [2] BFloat16 Specialization (Little Endian ppc64le)
|
||||
// On ppc64le (LE): BF16 bits should land in the HIGH 16 bits of each float32.
|
||||
// Byte layout of float32 on LE: [byte0(LSB), byte1, byte2, byte3(MSB)]
|
||||
// We need BF16 in bytes2-3 (high half) with bytes0-1 zeroed.
|
||||
// vec_mergeh on LE interleaves elements 0..3: result_i = {a[i], b[i]}
|
||||
// So vec_mergeh(zeros_u16, raw_u16) gives for each uint16 pair:
|
||||
// uint16[2i] = zeros[i] -> low 16 bits of uint32 -> zeroed mantissa LSBs
|
||||
// uint16[2i+1] = raw[i] -> high 16 bits of uint32 -> BF16 bits
|
||||
// Cast to float32 gives exactly (bf16_bits << 16) per element.
|
||||
template <>
|
||||
FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
|
||||
__vector float& b0,
|
||||
__vector float& b1) {
|
||||
__vector unsigned short raw = vec_xl(
|
||||
0, reinterpret_cast<unsigned short*>(const_cast<c10::BFloat16*>(p)));
|
||||
__vector unsigned short zeros = vec_splat_u16(0);
|
||||
|
||||
// LE: zeros in low 16 bits, raw in high 16 bits → bf16 << 16 == float32
|
||||
b0 = (__vector float)vec_mergeh(zeros, raw);
|
||||
b1 = (__vector float)vec_mergel(zeros, raw);
|
||||
}
|
||||
|
||||
// Note: c10::Half (FP16) is not supported on PowerPC architecture
|
||||
|
||||
template <int32_t M, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4(
|
||||
const float* __restrict A, // [M x K]
|
||||
const kv_cache_t* __restrict B, // [K x 8]
|
||||
float* __restrict C, // [M x 8]
|
||||
int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
|
||||
static_assert(1 <= M && M <= 8, "M must be in [1,8]");
|
||||
|
||||
#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
|
||||
#define IF_M(i) if constexpr (M > (i))
|
||||
|
||||
// 1. Define A pointers
|
||||
#define DECL_A(i) const float* a##i = A + (i) * lda;
|
||||
ROWS_APPLY(DECL_A)
|
||||
#undef DECL_A
|
||||
|
||||
// 2. Define Accumulators (2 vectors covers 8 columns)
|
||||
#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1;
|
||||
ROWS_APPLY(DECL_ACC)
|
||||
#undef DECL_ACC
|
||||
|
||||
// 3. Initialize Accumulators (Load C or Zero)
|
||||
#define INIT_ACC(i) \
|
||||
IF_M(i) { \
|
||||
if (accumulate) { \
|
||||
acc##i##_0 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 0)); \
|
||||
acc##i##_1 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 4)); \
|
||||
} else { \
|
||||
acc##i##_0 = vec_splats(0.0f); \
|
||||
acc##i##_1 = vec_splats(0.0f); \
|
||||
} \
|
||||
}
|
||||
ROWS_APPLY(INIT_ACC)
|
||||
#undef INIT_ACC
|
||||
|
||||
int32_t k = 0;
|
||||
|
||||
for (; k + 3 < K; k += 4) {
|
||||
// Load 4 values of A for each Row M: A[k...k+3]
|
||||
#define LOAD_A4(i) \
|
||||
__vector float a##i##v; \
|
||||
IF_M(i) a##i##v = vec_xl(0, const_cast<float*>(a##i + k));
|
||||
ROWS_APPLY(LOAD_A4)
|
||||
#undef LOAD_A4
|
||||
|
||||
// FMA for specific lane L of A
|
||||
// ppc64le: vec_madd(b, vec_splat(a, lane), acc)
|
||||
#define FMAS_LANE(i, aiv, L) \
|
||||
IF_M(i) { \
|
||||
__vector float a_broad = vec_splat(aiv, L); \
|
||||
acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \
|
||||
acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \
|
||||
}
|
||||
|
||||
// Unroll K=0..3
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
|
||||
#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
|
||||
ROWS_APPLY(STEP_K0)
|
||||
#undef STEP_K0
|
||||
}
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
|
||||
#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
|
||||
ROWS_APPLY(STEP_K1)
|
||||
#undef STEP_K1
|
||||
}
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
|
||||
#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
|
||||
ROWS_APPLY(STEP_K2)
|
||||
#undef STEP_K2
|
||||
}
|
||||
{
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
|
||||
#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
|
||||
ROWS_APPLY(STEP_K3)
|
||||
#undef STEP_K3
|
||||
}
|
||||
#undef FMAS_LANE
|
||||
}
|
||||
|
||||
for (; k < K; ++k) {
|
||||
__vector float b0, b1;
|
||||
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
|
||||
#define TAIL_ROW(i) \
|
||||
IF_M(i) { \
|
||||
__vector float ai = vec_splats(*(a##i + k)); \
|
||||
acc##i##_0 = vec_madd(b0, ai, acc##i##_0); \
|
||||
acc##i##_1 = vec_madd(b1, ai, acc##i##_1); \
|
||||
}
|
||||
ROWS_APPLY(TAIL_ROW)
|
||||
#undef TAIL_ROW
|
||||
}
|
||||
|
||||
#define STORE_ROW(i) \
|
||||
IF_M(i) { \
|
||||
vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \
|
||||
vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \
|
||||
}
|
||||
ROWS_APPLY(STORE_ROW)
|
||||
#undef STORE_ROW
|
||||
|
||||
#undef ROWS_APPLY
|
||||
#undef IF_M
|
||||
}
|
||||
|
||||
template <int32_t N, typename kv_cache_t>
|
||||
FORCE_INLINE void gemm_macro_ppc64le_Mx8_Ku4(const float* __restrict A,
|
||||
const kv_cache_t* __restrict B,
|
||||
float* __restrict C, int32_t M,
|
||||
int32_t K, int64_t lda,
|
||||
int64_t ldb, int64_t ldc,
|
||||
bool accumulate) {
|
||||
static_assert(N % 8 == 0, "N must be a multiple of 8");
|
||||
for (int32_t m = 0; m < M;) {
|
||||
int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
|
||||
const float* Ab = A + m * lda;
|
||||
float* Cb = C + m * ldc;
|
||||
|
||||
for (int32_t n = 0; n < N; n += 8) {
|
||||
const kv_cache_t* Bn = B + n;
|
||||
float* Cn = Cb + n;
|
||||
switch (mb) {
|
||||
case 8:
|
||||
gemm_micro_ppc64le_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
case 4:
|
||||
gemm_micro_ppc64le_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
case 2:
|
||||
gemm_micro_ppc64le_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
default:
|
||||
gemm_micro_ppc64le_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
|
||||
K, accumulate);
|
||||
break;
|
||||
}
|
||||
}
|
||||
m += mb;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename kv_cache_t>
|
||||
class TileGemmPPC64 {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
float* __restrict__ a_tile,
|
||||
kv_cache_t* __restrict__ b_tile,
|
||||
float* __restrict__ c_tile, const int64_t lda,
|
||||
const int64_t ldb, const int64_t ldc,
|
||||
const int32_t block_size,
|
||||
const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
gemm_macro_ppc64le_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
|
||||
} else {
|
||||
gemm_macro_ppc64le_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
|
||||
a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
|
||||
accum_c);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t, int64_t head_dim>
|
||||
class AttentionImpl<ISA::VSX, scalar_t, head_dim> {
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
using kv_cache_t = scalar_t;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = float;
|
||||
|
||||
constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
|
||||
constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
|
||||
constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
|
||||
constexpr static int64_t HeadDim = head_dim;
|
||||
constexpr static ISA ISAType = ISA::VSX;
|
||||
constexpr static bool scale_on_logits =
|
||||
false; // Scale is applied to Q during copy
|
||||
|
||||
public:
|
||||
AttentionImpl() {}
|
||||
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
attention<TileGemmPPC64<kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
|
||||
// Strides for Memory Layout
|
||||
constexpr static int64_t k_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return BlockSizeAlignment; // [head_dim, block_size] layout
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_token_group_stride(
|
||||
const int32_t block_size) {
|
||||
return head_dim * BlockSizeAlignment;
|
||||
}
|
||||
|
||||
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
|
||||
return HeadDimAlignment;
|
||||
}
|
||||
|
||||
static void copy_q_heads_tile(scalar_t* __restrict__ src,
|
||||
float* __restrict__ q_buffer,
|
||||
const int32_t q_num,
|
||||
const int32_t q_heads_per_kv,
|
||||
const int64_t q_num_stride,
|
||||
const int64_t q_head_stride, float scale) {
|
||||
__vector float scale_vec = vec_splats(scale);
|
||||
constexpr bool is_bf16 = std::is_same<scalar_t, c10::BFloat16>::value;
|
||||
|
||||
for (int32_t i = 0; i < q_num; ++i) {
|
||||
for (int32_t h = 0; h < q_heads_per_kv; ++h) {
|
||||
scalar_t* curr_src = src + i * q_num_stride + h * q_head_stride;
|
||||
float* curr_dst =
|
||||
q_buffer + i * q_heads_per_kv * head_dim + h * head_dim;
|
||||
|
||||
int32_t d = 0;
|
||||
for (; d <= head_dim - 8; d += 8) {
|
||||
__vector float v0, v1;
|
||||
load_row8_B_as_f32<scalar_t>(curr_src + d, v0, v1);
|
||||
|
||||
v0 = vec_mul(v0, scale_vec);
|
||||
v1 = vec_mul(v1, scale_vec);
|
||||
|
||||
vec_xst(v0, 0, curr_dst + d);
|
||||
vec_xst(v1, 0, curr_dst + d + 4);
|
||||
}
|
||||
|
||||
for (; d < head_dim; ++d) {
|
||||
float val = static_cast<float>(curr_src[d]);
|
||||
curr_dst[d] = val * scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void reshape_and_cache(
|
||||
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
|
||||
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
|
||||
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
|
||||
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
|
||||
const int64_t head_num, const int64_t key_head_num_stride,
|
||||
const int64_t value_head_num_stride, const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size, const int64_t block_size_stride,
|
||||
const float k_inv = 0.0f, const float v_inv = 0.0f) {
|
||||
// k_inv and v_inv are unused on VSX: FP8 KV cache is not supported on
|
||||
// PowerPC. The parameters are present to match the common interface.
|
||||
#pragma omp parallel for collapse(2)
|
||||
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
|
||||
const int64_t pos = slot_mapping[token_idx];
|
||||
if (pos < 0) continue;
|
||||
|
||||
const int64_t block_idx = pos / block_size;
|
||||
const int64_t block_offset = pos % block_size;
|
||||
|
||||
{
|
||||
const scalar_t* key_src = key + token_idx * key_token_num_stride +
|
||||
head_idx * key_head_num_stride;
|
||||
scalar_t* key_dst = key_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset;
|
||||
|
||||
for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
|
||||
key_dst[j] = key_src[i];
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
const scalar_t* val_src = value + token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
scalar_t* val_dst = value_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride +
|
||||
block_offset * head_dim;
|
||||
|
||||
std::memcpy(val_dst, val_src, sizeof(scalar_t) * head_dim);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_attention
|
||||
|
||||
#undef BLOCK_SIZE_ALIGNMENT
|
||||
#undef HEAD_SIZE_ALIGNMENT
|
||||
#undef MAX_Q_HEAD_NUM_PER_ITER
|
||||
|
||||
#endif // CPU_ATTN_VSX_HPP
|
||||
@@ -486,6 +486,10 @@ struct FP32Vec16 : public VectorizedRegWrapper<FP32Vec16, 4, float> {
|
||||
|
||||
explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
|
||||
|
||||
// FP8 stub: dead code on ARM (fp8 KV cache is x86-only), needed for
|
||||
// load_b_pair_vec template to compile on all platforms.
|
||||
explicit FP32Vec16(const BF16Vec32&, int) : Base() {}
|
||||
|
||||
explicit FP32Vec16(const FP16Vec16& v) {
|
||||
reg.val[0] = Vectorized<float>(vcvt_f32_f16(vget_low_f16(v.reg.val[0])));
|
||||
reg.val[1] = Vectorized<float>(vcvt_f32_f16(vget_high_f16(v.reg.val[0])));
|
||||
|
||||
@@ -6,6 +6,9 @@
|
||||
|
||||
namespace vec_op {
|
||||
|
||||
struct fp8_e4m3_tag {};
|
||||
struct fp8_e5m2_tag {};
|
||||
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
|
||||
@@ -145,6 +148,9 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
}
|
||||
|
||||
void save(void* ptr) const { *reinterpret_cast<f16x32_t*>(ptr) = reg; }
|
||||
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : reg{} {}
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : reg{} {}
|
||||
};
|
||||
|
||||
struct FP32Vec4 : public Vec<FP32Vec4> {
|
||||
@@ -302,6 +308,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
|
||||
FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {};
|
||||
|
||||
// FP8 stub: dead code on scalar path (fp8 KV cache is x86-only), needed for
|
||||
// load_b_pair_vec template to compile on all platforms.
|
||||
explicit FP32Vec16(const BF16Vec32&, int) : reg{} {}
|
||||
|
||||
FP32Vec16 operator*(const FP32Vec16& b) const {
|
||||
f32x16_t ret;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
|
||||
@@ -9,6 +9,10 @@
|
||||
|
||||
namespace vec_op {
|
||||
|
||||
// FP8 tag types for tag dispatch (see cpu_attn_vec.hpp)
|
||||
struct fp8_e4m3_tag {};
|
||||
struct fp8_e5m2_tag {};
|
||||
|
||||
// FIXME: FP16 is not fully supported in Torch-CPU
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
@@ -142,6 +146,9 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
: reg({vec8_data.reg, vec8_data.reg, vec8_data.reg, vec8_data.reg}) {}
|
||||
|
||||
void save(void* ptr) const { *reinterpret_cast<ss16x8x4_t*>(ptr) = reg; }
|
||||
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : reg{} {}
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : reg{} {}
|
||||
};
|
||||
|
||||
struct FP32Vec4 : public Vec<FP32Vec4> {
|
||||
@@ -404,6 +411,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
|
||||
explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
|
||||
|
||||
// FP8 stub: dead code on PowerPC (fp8 KV cache is x86-only), needed for
|
||||
// load_b_pair_vec template to compile on all platforms.
|
||||
explicit FP32Vec16(const BF16Vec32&, int) : reg{} {}
|
||||
|
||||
explicit FP32Vec16(const INT32Vec16& v) {
|
||||
reg.val[0] = vec_ctf(v.reg.val[0], 0);
|
||||
reg.val[1] = vec_ctf(v.reg.val[1], 0);
|
||||
|
||||
@@ -688,6 +688,10 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
|
||||
explicit FP32Vec16(const BF16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
|
||||
|
||||
// FP8 stub: dead code on s390x (fp8 KV cache is x86-only), needed for
|
||||
// load_b_pair_vec template to compile on all platforms.
|
||||
explicit FP32Vec16(const BF16Vec32&, int) : reg{} {}
|
||||
|
||||
FP32Vec16 operator*(const FP32Vec16& b) const {
|
||||
return FP32Vec16(f32x4x4_t({vec_mul(reg.val[0], b.reg.val[0]),
|
||||
vec_mul(reg.val[1], b.reg.val[1]),
|
||||
|
||||
+133
-6
@@ -122,9 +122,17 @@ struct FP16Vec16 : public Vec<FP16Vec16> {
|
||||
void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
|
||||
|
||||
void save(void* ptr, const int elem_num) const {
|
||||
#ifdef __AVX512BW__
|
||||
constexpr uint32_t M = 0xFFFFFFFF;
|
||||
__mmask16 mask = _cvtu32_mask16(M >> (32 - elem_num));
|
||||
_mm256_mask_storeu_epi16(ptr, mask, reg);
|
||||
#else
|
||||
// Fallback for lack of 16-bit masked store
|
||||
int16_t tmp[VEC_ELEM_NUM];
|
||||
_mm256_storeu_si256((__m256i*)tmp, reg);
|
||||
for (int i = 0; i < elem_num; ++i)
|
||||
reinterpret_cast<int16_t*>(ptr)[i] = tmp[i];
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
@@ -161,9 +169,17 @@ struct BF16Vec16 : public Vec<BF16Vec16> {
|
||||
void save(void* ptr) const { _mm256_storeu_si256((__m256i*)ptr, reg); }
|
||||
|
||||
void save(void* ptr, const int elem_num) const {
|
||||
#ifdef __AVX512BW__
|
||||
constexpr uint32_t M = 0xFFFFFFFF;
|
||||
__mmask16 mask = _cvtu32_mask16(M >> (32 - elem_num));
|
||||
_mm256_mask_storeu_epi16(ptr, mask, reg);
|
||||
#else
|
||||
// Fallback for lack of 16-bit masked store
|
||||
int16_t tmp[VEC_ELEM_NUM];
|
||||
_mm256_storeu_si256((__m256i*)tmp, reg);
|
||||
for (int i = 0; i < elem_num; ++i)
|
||||
reinterpret_cast<int16_t*>(ptr)[i] = tmp[i];
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
@@ -247,13 +263,12 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
explicit BF16Vec32(__m256i low, __m256i high)
|
||||
: reg_low(low), reg_high(high) {}
|
||||
|
||||
explicit BF16Vec32()
|
||||
: reg_low(_mm256_setzero_si256()), reg_high(_mm256_setzero_si256()) {}
|
||||
|
||||
explicit BF16Vec32(BF16Vec8& vec8_data)
|
||||
: reg_low((__m256i)_mm256_inserti32x4(
|
||||
_mm256_castsi128_si256((__m128i)vec8_data.reg),
|
||||
(__m128i)vec8_data.reg, 1)),
|
||||
reg_high((__m256i)_mm256_inserti32x4(
|
||||
_mm256_castsi128_si256((__m128i)vec8_data.reg),
|
||||
(__m128i)vec8_data.reg, 1)) {}
|
||||
: reg_low(_mm256_broadcastsi128_si256((__m128i)vec8_data.reg)),
|
||||
reg_high(_mm256_broadcastsi128_si256((__m128i)vec8_data.reg)) {}
|
||||
|
||||
// E4M3 decode (AVX2 path) — same bit-layout trick as the AVX512 variant
|
||||
// above. Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
|
||||
@@ -674,6 +689,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
_mm256_sub_ps(reg_high, b.reg_high));
|
||||
}
|
||||
|
||||
FP32Vec16 operator-() const {
|
||||
const __m256 neg = _mm256_set1_ps(-0.0f);
|
||||
return FP32Vec16(_mm256_xor_ps(reg_low, neg), _mm256_xor_ps(reg_high, neg));
|
||||
}
|
||||
|
||||
FP32Vec16 operator/(const FP32Vec16& b) const {
|
||||
return FP32Vec16(_mm256_div_ps(reg_low, b.reg_low),
|
||||
_mm256_div_ps(reg_high, b.reg_high));
|
||||
@@ -739,6 +759,85 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
_mm256_storeu_ps(ptr, reg_low);
|
||||
_mm256_storeu_ps(ptr + 8, reg_high);
|
||||
}
|
||||
|
||||
void save(float* ptr, const int elem_num) const {
|
||||
// Partial store: cmpgt produces a sign-bit mask (0xFFFFFFFF/0 per lane)
|
||||
// for the first elem_num lanes, applied across the two 8-wide halves.
|
||||
if (elem_num <= 8) {
|
||||
__m256i mask =
|
||||
_mm256_cmpgt_epi32(_mm256_set1_epi32(elem_num),
|
||||
_mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7));
|
||||
_mm256_maskstore_ps(ptr, mask, reg_low);
|
||||
} else {
|
||||
_mm256_storeu_ps(ptr, reg_low);
|
||||
__m256i mask =
|
||||
_mm256_cmpgt_epi32(_mm256_set1_epi32(elem_num - 8),
|
||||
_mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7));
|
||||
_mm256_maskstore_ps(ptr + 8, mask, reg_high);
|
||||
}
|
||||
}
|
||||
|
||||
FP32Vec16 clamp(const FP32Vec16& min, const FP32Vec16& max) const {
|
||||
return FP32Vec16(
|
||||
_mm256_min_ps(max.reg_low, _mm256_max_ps(min.reg_low, reg_low)),
|
||||
_mm256_min_ps(max.reg_high, _mm256_max_ps(min.reg_high, reg_high)));
|
||||
}
|
||||
|
||||
FP32Vec16 abs() const {
|
||||
const __m256 sign_mask = _mm256_set1_ps(-0.0f);
|
||||
return FP32Vec16(_mm256_andnot_ps(sign_mask, reg_low),
|
||||
_mm256_andnot_ps(sign_mask, reg_high));
|
||||
}
|
||||
|
||||
FP32Vec16 min(const FP32Vec16& b) const {
|
||||
return FP32Vec16(_mm256_min_ps(reg_low, b.reg_low),
|
||||
_mm256_min_ps(reg_high, b.reg_high));
|
||||
}
|
||||
|
||||
// Partial element-wise min over the first elem_num lanes only (tail path).
|
||||
// Scalar via AliasReg: AVX2 has no masked vminps, so we spill, loop, reload.
|
||||
FP32Vec16 min(const FP32Vec16& b, const int elem_num) const {
|
||||
AliasReg ar_this_low, ar_this_high, ar_b_low, ar_b_high;
|
||||
ar_this_low.reg = reg_low;
|
||||
ar_this_high.reg = reg_high;
|
||||
ar_b_low.reg = b.reg_low;
|
||||
ar_b_high.reg = b.reg_high;
|
||||
for (int i = 0; i < elem_num && i < 8; ++i)
|
||||
ar_this_low.values[i] =
|
||||
std::min(ar_this_low.values[i], ar_b_low.values[i]);
|
||||
for (int i = 0; i < elem_num - 8 && i < 8; ++i)
|
||||
ar_this_high.values[i] =
|
||||
std::min(ar_this_high.values[i], ar_b_high.values[i]);
|
||||
return FP32Vec16(ar_this_low.reg, ar_this_high.reg);
|
||||
}
|
||||
|
||||
// Partial element-wise max over the first elem_num lanes only (tail path).
|
||||
// Scalar via AliasReg: AVX2 has no masked vmaxps, so we spill, loop, reload.
|
||||
FP32Vec16 max(const FP32Vec16& b, const int elem_num) const {
|
||||
AliasReg ar_this_low, ar_this_high, ar_b_low, ar_b_high;
|
||||
ar_this_low.reg = reg_low;
|
||||
ar_this_high.reg = reg_high;
|
||||
ar_b_low.reg = b.reg_low;
|
||||
ar_b_high.reg = b.reg_high;
|
||||
for (int i = 0; i < elem_num && i < 8; ++i)
|
||||
ar_this_low.values[i] =
|
||||
std::max(ar_this_low.values[i], ar_b_low.values[i]);
|
||||
for (int i = 0; i < elem_num - 8 && i < 8; ++i)
|
||||
ar_this_high.values[i] =
|
||||
std::max(ar_this_high.values[i], ar_b_high.values[i]);
|
||||
return FP32Vec16(ar_this_low.reg, ar_this_high.reg);
|
||||
}
|
||||
|
||||
float reduce_min() const {
|
||||
__m256 v = _mm256_min_ps(reg_low, reg_high);
|
||||
__m256 v_shuffled = _mm256_permute_ps(v, 0b00001011);
|
||||
__m256 v_min = _mm256_min_ps(v, v_shuffled);
|
||||
v_shuffled = _mm256_permute_ps(v_min, 0b00000001);
|
||||
v_min = _mm256_min_ps(v_min, v_shuffled);
|
||||
v_shuffled = _mm256_permute2f128_ps(v_min, v_min, 0b00000001);
|
||||
v_min = _mm256_min_ps(v_min, v_shuffled);
|
||||
return _mm256_cvtss_f32(v_min);
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
@@ -791,6 +890,34 @@ struct INT8Vec64 : public Vec<INT8Vec64> {
|
||||
// non-temporal save
|
||||
void nt_save(int8_t* ptr) { _mm512_stream_si512((__m512i*)ptr, reg); }
|
||||
};
|
||||
#else
|
||||
struct INT8Vec16 : public Vec<INT8Vec16> {
|
||||
constexpr static int VEC_ELEM_NUM = 16;
|
||||
union AliasReg {
|
||||
__m128i reg;
|
||||
int8_t values[VEC_ELEM_NUM];
|
||||
};
|
||||
|
||||
__m128i reg;
|
||||
|
||||
explicit INT8Vec16(const FP32Vec16& vec) {
|
||||
__m256i lo_i32 = _mm256_cvtps_epi32(vec.reg_low);
|
||||
__m256i hi_i32 = _mm256_cvtps_epi32(vec.reg_high);
|
||||
__m256i packed16 = _mm256_packs_epi32(lo_i32, hi_i32);
|
||||
packed16 = _mm256_permute4x64_epi64(packed16, 0xD8);
|
||||
__m256i packed8 = _mm256_packs_epi16(packed16, _mm256_setzero_si256());
|
||||
packed8 = _mm256_permute4x64_epi64(packed8, 0xD8);
|
||||
reg = _mm256_castsi256_si128(packed8);
|
||||
}
|
||||
|
||||
void save(int8_t* ptr) const { _mm_storeu_si128((__m128i*)ptr, reg); }
|
||||
|
||||
void save(int8_t* ptr, const int elem_num) const {
|
||||
AliasReg ar;
|
||||
ar.reg = reg;
|
||||
for (int i = 0; i < elem_num; ++i) ptr[i] = ar.values[i];
|
||||
}
|
||||
};
|
||||
#endif
|
||||
|
||||
template <typename T>
|
||||
|
||||
@@ -215,7 +215,7 @@ void dynamic_quant_epilogue(const float* input, scalar_t* output,
|
||||
float zp_scale_val = a_scale[i] * static_cast<float>(azp[i]);
|
||||
token_zp_scale_vec = cvt_vec_t(zp_scale_val);
|
||||
}
|
||||
for (; j < hidden_size - vec_elem_num; ++j) {
|
||||
for (; j < hidden_size - vec_elem_num; j += vec_elem_num) {
|
||||
cvt_vec_t elems_fp32(input_ptr + j);
|
||||
elems_fp32 = elems_fp32 * token_scale_vec;
|
||||
if constexpr (AZP) {
|
||||
|
||||
@@ -20,6 +20,7 @@ ISA_TYPES = {
|
||||
"VEC16": 2,
|
||||
"NEON": 3,
|
||||
"VXE": 4,
|
||||
"VSX": 5,
|
||||
}
|
||||
|
||||
# KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
|
||||
@@ -37,7 +38,7 @@ KV_CACHE_CPP_TYPES = {
|
||||
}
|
||||
|
||||
# ISAs supported for head_dims divisible by 32
|
||||
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE"]
|
||||
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "VSX"]
|
||||
|
||||
# ISAs supported for head_dims divisible by 16 only
|
||||
ISA_FOR_16 = ["VEC16"]
|
||||
@@ -148,6 +149,10 @@ def generate_header_file() -> str:
|
||||
#include "cpu_attn_vxe.hpp"
|
||||
#endif
|
||||
|
||||
#ifdef __powerpc__
|
||||
#include "cpu_attn_vsx.hpp"
|
||||
#endif
|
||||
|
||||
"""
|
||||
|
||||
header += generate_helper_function()
|
||||
@@ -207,6 +212,11 @@ def generate_header_file() -> str:
|
||||
["VXE", "VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__powerpc__)",
|
||||
["VSX", "VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__AVX512F__)",
|
||||
["VEC", "VEC16"],
|
||||
@@ -223,7 +233,8 @@ def generate_header_file() -> str:
|
||||
fp8=False,
|
||||
)
|
||||
header += (
|
||||
"#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */\n\n"
|
||||
"#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / "
|
||||
"__s390x__ / __powerpc__ */\n\n"
|
||||
"#endif // CPU_ATTN_DISPATCH_GENERATED_H\n"
|
||||
)
|
||||
|
||||
|
||||
@@ -447,7 +447,7 @@ INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
|
||||
INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
|
||||
|
||||
at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2,
|
||||
std::vector<int64_t> block_size, std::optional<at::Tensor>& bias,
|
||||
std::vector<int64_t> block_size, const std::optional<at::Tensor>& bias,
|
||||
at::ScalarType out_dtype, bool is_vnni) {
|
||||
RECORD_FUNCTION("sgl-kernel::fp8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales2, block_size, bias}));
|
||||
|
||||
|
||||
@@ -77,6 +77,13 @@ at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
at::ScalarType out_dtype, bool is_vnni);
|
||||
|
||||
// Adapted from sglang: FP8 W8A16 kernel
|
||||
at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
|
||||
at::Tensor& scales2,
|
||||
std::vector<int64_t> block_size,
|
||||
const std::optional<at::Tensor>& bias,
|
||||
at::ScalarType out_dtype, bool is_vnni);
|
||||
|
||||
// Adapted from sglang: INT4 W4A8 kernels
|
||||
std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
|
||||
at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
|
||||
@@ -271,8 +278,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
|
||||
|
||||
// Quantization
|
||||
#if defined(__AVX512F__) || (defined(__aarch64__) && !defined(__APPLE__)) || \
|
||||
defined(__powerpc64__)
|
||||
#if defined(__AVX512F__) || defined(__AVX2__) || \
|
||||
(defined(__aarch64__) && !defined(__APPLE__)) || defined(__powerpc64__)
|
||||
// Helper function to release oneDNN handlers
|
||||
ops.def("release_dnnl_matmul_handler(int handler) -> ()",
|
||||
&release_dnnl_matmul_handler);
|
||||
@@ -376,6 +383,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"int4_scaled_mm_cpu(Tensor(a0!) x, Tensor(a1!) w, Tensor(a2!) w_zeros, "
|
||||
"Tensor(a3!) w_scales, Tensor? bias) -> Tensor");
|
||||
ops.impl("int4_scaled_mm_cpu", torch::kCPU, &int4_scaled_mm_cpu);
|
||||
|
||||
// Adapted from sglang: FP8 W8A16 kernel
|
||||
ops.def(
|
||||
"fp8_scaled_mm_cpu(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!) "
|
||||
"scales2, SymInt[] block_size, Tensor? bias, ScalarType out_dtype, "
|
||||
"bool is_vnni) -> Tensor");
|
||||
ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
|
||||
#endif
|
||||
|
||||
// CPU attention kernels
|
||||
|
||||
+1
-1
@@ -54,7 +54,7 @@ struct Counter {
|
||||
};
|
||||
|
||||
inline int64_t get_available_l2_size() {
|
||||
#if defined(__s390x__)
|
||||
#if defined(__s390x__) || defined(__powerpc__)
|
||||
static int64_t size = []() {
|
||||
uint32_t l2_cache_size = 0;
|
||||
auto caps = at::cpu::get_cpu_capabilities();
|
||||
|
||||
@@ -232,28 +232,6 @@ void unmap_and_release(unsigned long long device, ssize_t size,
|
||||
}
|
||||
}
|
||||
|
||||
// ROCm workaround: hipMemRelease does not return physical VRAM to the
|
||||
// free pool while the virtual-address reservation is still held.
|
||||
// Cycling cuMemAddressFree → cuMemAddressReserve (at the same address)
|
||||
// forces the driver to actually release the physical pages while keeping
|
||||
// the same VA available for a later create_and_map.
|
||||
if (first_error == no_error) {
|
||||
first_error = cuMemAddressFree(d_mem, size);
|
||||
if (first_error == no_error) {
|
||||
CUdeviceptr d_mem_new = 0;
|
||||
first_error = cuMemAddressReserve(&d_mem_new, size, 0, d_mem, 0);
|
||||
if (first_error == no_error && d_mem_new != d_mem) {
|
||||
cuMemAddressFree(d_mem_new, size);
|
||||
snprintf(error_msg, sizeof(error_msg),
|
||||
"ROCm: VA re-reserve got %p instead of %p", (void*)d_mem_new,
|
||||
(void*)d_mem);
|
||||
error_code = CUresult(1);
|
||||
std::cerr << error_msg << std::endl;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (first_error != no_error) {
|
||||
CUDA_CHECK(first_error);
|
||||
}
|
||||
|
||||
@@ -29,7 +29,11 @@
|
||||
*/
|
||||
|
||||
#include <cmath>
|
||||
#include <cuda_fp8.h>
|
||||
#ifndef USE_ROCM
|
||||
#include <cuda_fp8.h>
|
||||
#else
|
||||
#include <hip/hip_fp8.h>
|
||||
#endif
|
||||
#include <cuda_runtime.h>
|
||||
#include <type_traits>
|
||||
|
||||
@@ -42,7 +46,23 @@
|
||||
#include "type_convert.cuh"
|
||||
|
||||
#ifndef FINAL_MASK
|
||||
#define FINAL_MASK 0xffffffffu
|
||||
#ifdef USE_ROCM
|
||||
#define FINAL_MASK 0xffffffffffffffffULL
|
||||
#else
|
||||
#define FINAL_MASK 0xffffffffu
|
||||
#endif
|
||||
#endif
|
||||
|
||||
#ifdef USE_ROCM
|
||||
// ROCm-compatible FP8 conversion helpers
|
||||
__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
|
||||
#if defined(HIP_FP8_TYPE_OCP)
|
||||
__hip_fp8_e4m3 fp8_val(val);
|
||||
#else
|
||||
__hip_fp8_e4m3_fnuz fp8_val(val);
|
||||
#endif
|
||||
return reinterpret_cast<uint8_t&>(fp8_val);
|
||||
}
|
||||
#endif
|
||||
|
||||
namespace vllm {
|
||||
@@ -314,9 +334,13 @@ __global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
float scaled = elements[i] * inv_scale;
|
||||
scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
|
||||
#ifndef USE_ROCM
|
||||
__nv_fp8_storage_t s =
|
||||
__nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
|
||||
out_bytes[i] = static_cast<uint8_t>(s);
|
||||
#else
|
||||
out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
|
||||
#endif
|
||||
}
|
||||
// One 16-byte STG per lane.
|
||||
*reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
|
||||
@@ -384,6 +408,7 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
|
||||
// PDL: enable programmatic stream serialization whenever the hardware
|
||||
// supports it (SM90+). On pre-Hopper GPUs the attribute is unavailable,
|
||||
// so leave numAttrs = 0 and launch as a regular kernel.
|
||||
#ifndef USE_ROCM
|
||||
static int const sm_version = getSMVersion();
|
||||
// Host-side guard: the device kernel body is compiled as a no-op for
|
||||
// bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
|
||||
@@ -410,6 +435,15 @@ void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
|
||||
q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
|
||||
num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
|
||||
kv_block_stride);
|
||||
#else
|
||||
// ROCm: use standard kernel launch syntax (no PDL/stream serialization)
|
||||
// clang-format off
|
||||
fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>
|
||||
<<<grid, kBlockSize, 0, stream>>>(
|
||||
q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
|
||||
eps, num_tokens_full, num_tokens_insert, num_heads_q,
|
||||
cache_block_size, kv_block_stride);
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace deepseek_v4_fused_ops
|
||||
|
||||
@@ -236,17 +236,41 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
|
||||
#undef LAUNCH_KERNEL
|
||||
}
|
||||
|
||||
template <typename T, typename DST_DTYPE>
|
||||
__global__ void per_token_group_quant_8bit_packed_kernel(
|
||||
// Register-resident fast path for group_size==128.
|
||||
//
|
||||
// Each thread holds 16 source elements (32 B = uint4 x 2) in registers across
|
||||
// the absmax reduce -> scale compute -> quantize pipeline. No shared memory.
|
||||
// UE8M0 scale extracted via bit math (bit-exact with exp2f(ceilf(log2f))).
|
||||
//
|
||||
// Loads two contiguous uint4s (16 B + 16 B = 32 B) per thread; on Blackwell
|
||||
// nvcc fuses these into a single 256-bit LDG.E.256.
|
||||
//
|
||||
// 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>
|
||||
__global__ void per_token_group_quant_8bit_packed_register_kernel(
|
||||
const T* __restrict__ input, void* __restrict__ output_q,
|
||||
unsigned int* __restrict__ output_s_packed, const int group_size,
|
||||
const int num_groups_padded, const int groups_per_block,
|
||||
const int padded_groups_per_row, const int groups_per_row, const int mn,
|
||||
const int tma_aligned_mn, const int num_scale_elems, const float eps,
|
||||
unsigned int* __restrict__ output_s_packed, const int64_t num_groups_padded,
|
||||
const int groups_per_block, const int padded_groups_per_row,
|
||||
const int groups_per_row, const int mn, const int output_q_mn_extent,
|
||||
const int tma_aligned_mn, const int64_t num_scale_elems, const float eps,
|
||||
const float min_8bit, const float max_8bit) {
|
||||
const int threads_per_group = 16;
|
||||
const int64_t local_group_id = threadIdx.x / threads_per_group;
|
||||
const int lane_id = threadIdx.x % threads_per_group;
|
||||
static_assert(GROUP_SIZE == 128, "fast path supports GROUP_SIZE==128");
|
||||
constexpr int THREADS_PER_GROUP = 8;
|
||||
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");
|
||||
|
||||
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;
|
||||
@@ -254,141 +278,207 @@ __global__ void per_token_group_quant_8bit_packed_kernel(
|
||||
return;
|
||||
}
|
||||
|
||||
// map flat group id to 2D indices (mn_idx, sf_k_idx)
|
||||
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);
|
||||
|
||||
// whether it is a valid group (not padding)
|
||||
const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
|
||||
|
||||
// shared memory to cache each group's data to avoid double DRAM reads.
|
||||
extern __shared__ __align__(16) char smem_raw[];
|
||||
T* smem = reinterpret_cast<T*>(smem_raw);
|
||||
T* smem_group = smem + local_group_id * group_size;
|
||||
|
||||
// compute scale for valid groups
|
||||
float y_s = 0.f;
|
||||
// Load 16 input elements (32 B) into registers as two adjacent uint4
|
||||
// loads. nvcc keeps these as 2x LDG.E.128 on sm_100; the per-thread cost
|
||||
// is dominated by HBM bandwidth at large MN, so a fused 256-bit load via
|
||||
// inline PTX gave no measurable speedup.
|
||||
// alignas(16) is required so the uint4* reinterpret_cast below is
|
||||
// well-defined for T == bf16/fp16 (default alignof is 2).
|
||||
alignas(16) T regs[VEC_SIZE];
|
||||
float local_absmax = eps;
|
||||
if (is_valid_group) {
|
||||
const T* group_input =
|
||||
input + static_cast<int64_t>(mn_idx) * groups_per_row * group_size +
|
||||
sf_k_idx * group_size;
|
||||
y_s = ComputeGroupScale<T, true>(group_input, smem_group, group_size,
|
||||
lane_id, threads_per_group, eps, max_8bit);
|
||||
input + static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
|
||||
sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
|
||||
uint4* dst = reinterpret_cast<uint4*>(®s[0]);
|
||||
const uint4* src = reinterpret_cast<const uint4*>(group_input);
|
||||
dst[0] = src[0];
|
||||
dst[1] = src[1];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < VEC_SIZE; ++i) {
|
||||
float v = fabsf(static_cast<float>(regs[i]));
|
||||
local_absmax = fmaxf(local_absmax, v);
|
||||
}
|
||||
}
|
||||
|
||||
// pack 4 scales into a uint32 exponent
|
||||
// 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the
|
||||
// 8 lanes within the warp that share a group.
|
||||
unsigned mask = 0xffu << (threadIdx.x & 24u);
|
||||
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4));
|
||||
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2));
|
||||
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1));
|
||||
|
||||
float y_s = local_absmax / max_8bit;
|
||||
y_s = fmaxf(y_s, 1e-10f);
|
||||
uint32_t bits = __float_as_uint(y_s);
|
||||
uint32_t exp_bits = (bits >> 23) & 0xffu;
|
||||
uint32_t mant_bits = bits & 0x7fffffu;
|
||||
uint8_t exp_byte =
|
||||
static_cast<uint8_t>(exp_bits + (mant_bits != 0u ? 1u : 0u));
|
||||
|
||||
// Lane 0 writes the packed scale byte.
|
||||
if (lane_id == 0) {
|
||||
// each uint32 in output_s_packed stores 4 packed scales
|
||||
const int sf_k_pack_idx = sf_k_idx / 4;
|
||||
const int pos = sf_k_idx % 4;
|
||||
const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx;
|
||||
|
||||
if (is_valid_group) {
|
||||
// reinterpret the UE8M0 scale y_s as IEEE bits, extract the 8-bit
|
||||
// exponent, and place it into the correct byte of the 32-bit word.
|
||||
const unsigned int bits = __float_as_uint(y_s);
|
||||
const uint8_t exponent = static_cast<uint8_t>((bits >> 23u) & 0xffu);
|
||||
reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = exponent;
|
||||
reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = exp_byte;
|
||||
} else if (out_idx < num_scale_elems) {
|
||||
// write zero for padding groups if within bounds of output_s_packed
|
||||
reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (is_valid_group) {
|
||||
DST_DTYPE* group_output =
|
||||
static_cast<DST_DTYPE*>(output_q) +
|
||||
static_cast<int64_t>(mn_idx) * groups_per_row * group_size +
|
||||
sf_k_idx * group_size;
|
||||
QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
|
||||
threads_per_group, y_s, min_8bit, max_8bit);
|
||||
// For padded mn rows that fall within output_q's allocated extent, write
|
||||
// a uint4 of zeros to keep the buffer clean for downstream TMA loads.
|
||||
// Skip writes for sf_k padding (those positions don't exist in output_q).
|
||||
if (!is_valid_group) {
|
||||
if (sf_k_idx < groups_per_row && mn_idx >= mn &&
|
||||
mn_idx < output_q_mn_extent) {
|
||||
DST_DTYPE* group_output =
|
||||
static_cast<DST_DTYPE*>(output_q) +
|
||||
static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
|
||||
sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
|
||||
*reinterpret_cast<uint4*>(group_output) = make_uint4(0, 0, 0, 0);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// Reconstruct y_s as a power-of-2 float and use its reciprocal.
|
||||
float y_s_q = __uint_as_float(static_cast<uint32_t>(exp_byte) << 23);
|
||||
float inv_y = 1.0f / y_s_q;
|
||||
|
||||
// Quantize and pack into 16 fp8/int8 bytes (= uint4). VEC_SIZE==16 so we
|
||||
// fill four 32-bit words, four bytes each.
|
||||
uint32_t packed_lo = 0;
|
||||
uint32_t packed_lo_hi = 0;
|
||||
uint32_t packed_hi_lo = 0;
|
||||
uint32_t packed_hi = 0;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < VEC_SIZE; ++i) {
|
||||
float q =
|
||||
fminf(fmaxf(static_cast<float>(regs[i]) * inv_y, min_8bit), max_8bit);
|
||||
DST_DTYPE qb = DST_DTYPE(q);
|
||||
uint8_t byte = *reinterpret_cast<uint8_t*>(&qb);
|
||||
const int shift = (i & 3) * 8;
|
||||
if (i < 4) {
|
||||
packed_lo |= static_cast<uint32_t>(byte) << shift;
|
||||
} else if (i < 8) {
|
||||
packed_lo_hi |= static_cast<uint32_t>(byte) << shift;
|
||||
} else if (i < 12) {
|
||||
packed_hi_lo |= static_cast<uint32_t>(byte) << shift;
|
||||
} else {
|
||||
packed_hi |= static_cast<uint32_t>(byte) << shift;
|
||||
}
|
||||
}
|
||||
|
||||
uint4 packed_out =
|
||||
make_uint4(packed_lo, packed_lo_hi, packed_hi_lo, packed_hi);
|
||||
DST_DTYPE* group_output =
|
||||
static_cast<DST_DTYPE*>(output_q) +
|
||||
static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
|
||||
sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
|
||||
*reinterpret_cast<uint4*>(group_output) = packed_out;
|
||||
}
|
||||
|
||||
// Public entry point: register-resident packed quant kernel.
|
||||
// Constraints: group_size == 128 and bf16/fp16 input.
|
||||
void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
|
||||
torch::stable::Tensor& output_q,
|
||||
torch::stable::Tensor& output_s_packed,
|
||||
int64_t group_size, double eps,
|
||||
double min_8bit, double max_8bit) {
|
||||
STD_TORCH_CHECK(group_size == 128,
|
||||
"per_token_group_quant_8bit_packed only supports "
|
||||
"group_size==128, got ",
|
||||
group_size, ".");
|
||||
const auto in_dtype = input.scalar_type();
|
||||
STD_TORCH_CHECK(
|
||||
in_dtype == torch::headeronly::ScalarType::Half ||
|
||||
in_dtype == torch::headeronly::ScalarType::BFloat16,
|
||||
"per_token_group_quant_8bit_packed only supports bf16/fp16 input.");
|
||||
|
||||
STD_TORCH_CHECK(input.is_contiguous());
|
||||
STD_TORCH_CHECK(output_q.is_contiguous());
|
||||
|
||||
const int64_t k = input.size(-1);
|
||||
STD_TORCH_CHECK(k % group_size == 0, "Last dimension (", k,
|
||||
") must be divisible by group_size (", group_size, ").");
|
||||
STD_TORCH_CHECK(k % group_size == 0, "input last dim k=", k,
|
||||
" is not divisible by group_size=", group_size, ".");
|
||||
|
||||
const int64_t mn = input.numel() / k;
|
||||
const int64_t groups_per_row = k / group_size;
|
||||
|
||||
STD_TORCH_CHECK(output_s_packed.dim() == 2,
|
||||
"output_s_packed must be 2D, got dim=", output_s_packed.dim(),
|
||||
".");
|
||||
|
||||
const int64_t k_num_packed_sfk = (groups_per_row + 3) / 4;
|
||||
const int64_t tma_aligned_mn = ((mn + 3) / 4) * 4;
|
||||
|
||||
// output_q may be allocated with extra padded mn rows (e.g.,
|
||||
// (tma_aligned_mn, k)) so the kernel can zero-fill them in-line and the
|
||||
// caller can use torch.empty instead of torch.zeros. The grid only covers
|
||||
// up to tma_aligned_mn, so we cap the extent there.
|
||||
const int64_t output_q_mn_actual = output_q.numel() / k;
|
||||
STD_TORCH_CHECK(output_q_mn_actual >= mn,
|
||||
"output_q must have at least mn rows; got ",
|
||||
output_q_mn_actual, " rows for mn=", mn, ".");
|
||||
const int64_t output_q_mn_extent =
|
||||
output_q_mn_actual < tma_aligned_mn ? output_q_mn_actual : tma_aligned_mn;
|
||||
|
||||
STD_TORCH_CHECK(
|
||||
output_s_packed.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"output_s_packed must have dtype int32 for UE8M0-packed scales.");
|
||||
// DeepGEMM expects SFA scales in MN-major form with shape
|
||||
// [mn, ceil_div(K, 128 * 4)] and TMA-aligned stride on the last
|
||||
// dimension.
|
||||
"output_s_packed must be int32 for UE8M0-packed scales.");
|
||||
STD_TORCH_CHECK(output_s_packed.size(0) == mn &&
|
||||
output_s_packed.size(1) == k_num_packed_sfk,
|
||||
"output_s_packed shape must be [", mn, ", ", k_num_packed_sfk,
|
||||
"], but got [", output_s_packed.size(0), ", ",
|
||||
"]; got [", output_s_packed.size(0), ", ",
|
||||
output_s_packed.size(1), "].");
|
||||
// Verify column-major TMA-aligned layout
|
||||
STD_TORCH_CHECK(output_s_packed.stride(0) == 1 &&
|
||||
output_s_packed.stride(1) == tma_aligned_mn,
|
||||
"output_s_packed must have strides [1, ", tma_aligned_mn,
|
||||
"], but got [", output_s_packed.stride(0), ", ",
|
||||
"output_s_packed strides must be [1, ", tma_aligned_mn,
|
||||
"]; got [", output_s_packed.stride(0), ", ",
|
||||
output_s_packed.stride(1), "].");
|
||||
|
||||
cudaStream_t stream = get_current_cuda_stream();
|
||||
|
||||
constexpr int THREADS_PER_GROUP = 16;
|
||||
|
||||
// Expand the grid to cover MN and K padding so every byte in
|
||||
// output_s_packed is written (padding bytes get zeroed by the kernel).
|
||||
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;
|
||||
// Number of elements in output_s_packed.
|
||||
const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn;
|
||||
|
||||
const int groups_per_block = GetGroupsPerBlock(num_groups_padded);
|
||||
|
||||
auto dst_type = output_q.scalar_type();
|
||||
const int num_blocks = num_groups_padded / groups_per_block;
|
||||
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_PACKED_KERNEL(T, DST_DTYPE) \
|
||||
do { \
|
||||
dim3 grid(num_blocks); \
|
||||
dim3 block(num_threads); \
|
||||
size_t smem_bytes = \
|
||||
static_cast<size_t>(groups_per_block) * group_size * sizeof(T); \
|
||||
per_token_group_quant_8bit_packed_kernel<T, DST_DTYPE> \
|
||||
<<<grid, block, smem_bytes, stream>>>( \
|
||||
static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
|
||||
reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
|
||||
static_cast<int>(group_size), static_cast<int>(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>(tma_aligned_mn), \
|
||||
static_cast<int>(num_scale_elems), static_cast<float>(eps), \
|
||||
static_cast<float>(min_8bit), static_cast<float>(max_8bit)); \
|
||||
#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
|
||||
do { \
|
||||
dim3 grid(static_cast<unsigned int>(num_blocks)); \
|
||||
dim3 block(num_threads); \
|
||||
per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128> \
|
||||
<<<grid, block, 0, stream>>>( \
|
||||
static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
|
||||
reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
|
||||
num_groups_padded, groups_per_block, \
|
||||
static_cast<int>(padded_groups_per_row), \
|
||||
static_cast<int>(groups_per_row), static_cast<int>(mn), \
|
||||
static_cast<int>(output_q_mn_extent), \
|
||||
static_cast<int>(tma_aligned_mn), num_scale_elems, \
|
||||
static_cast<float>(eps), static_cast<float>(min_8bit), \
|
||||
static_cast<float>(max_8bit)); \
|
||||
} while (0)
|
||||
|
||||
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
|
||||
input.scalar_type(), "per_token_group_quant_8bit_packed", ([&] {
|
||||
VLLM_STABLE_DISPATCH_HALF_TYPES(
|
||||
input.scalar_type(), "per_token_group_quant_8bit_packed_register", ([&] {
|
||||
if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
|
||||
LAUNCH_PACKED_KERNEL(scalar_t, __nv_fp8_e4m3);
|
||||
LAUNCH_REG_KERNEL(scalar_t, __nv_fp8_e4m3);
|
||||
} else if (dst_type == torch::headeronly::ScalarType::Char) {
|
||||
LAUNCH_PACKED_KERNEL(scalar_t, int8_t);
|
||||
LAUNCH_REG_KERNEL(scalar_t, int8_t);
|
||||
} else {
|
||||
STD_TORCH_CHECK(
|
||||
false,
|
||||
@@ -397,7 +487,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
|
||||
}
|
||||
}));
|
||||
|
||||
#undef LAUNCH_PACKED_KERNEL
|
||||
#undef LAUNCH_REG_KERNEL
|
||||
}
|
||||
|
||||
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
|
||||
|
||||
@@ -8,3 +8,13 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
|
||||
torch::stable::Tensor& output_s,
|
||||
int64_t group_size, double eps, double min_8bit,
|
||||
double max_8bit, bool scale_ue8m0 = false);
|
||||
|
||||
// Public op: register-resident packed quant for the DeepGEMM Blackwell path.
|
||||
// Restricted to group_size == 128 and bf16/fp16 input; other configurations
|
||||
// raise STD_TORCH_CHECK. The legacy shared-memory fallback was removed because
|
||||
// no production caller (deep_gemm_moe / input_quant_fp8) uses other shapes.
|
||||
void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
|
||||
torch::stable::Tensor& output_q,
|
||||
torch::stable::Tensor& output_s_packed,
|
||||
int64_t group_size, double eps,
|
||||
double min_8bit, double max_8bit);
|
||||
|
||||
@@ -67,10 +67,6 @@ void shuffle_rows(const torch::Tensor& input_tensor,
|
||||
torch::Tensor& output_tensor);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16)
|
||||
torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
|
||||
torch::Tensor const& weight);
|
||||
|
||||
// DeepSeek V3 optimized router GEMM kernel for SM90+
|
||||
// Computes output = mat_a @ mat_b.T where:
|
||||
// mat_a: [num_tokens, hidden_dim] in bf16
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
// bf16 x bf16 -> fp32 router GEMM via cuBLAS.
|
||||
// Uses CUBLAS_COMPUTE_32F so bf16 operands accumulate into fp32,
|
||||
// matching TRT-LLM's cuBLAS fallback behaviour in dsv3RouterGemmOp.
|
||||
|
||||
#include <torch/all.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cublas_v2.h>
|
||||
|
||||
// cuBLAS column-major math for row-major PyTorch tensors:
|
||||
// weight[N,K]_row lda=K -> cuBLAS sees (K,N) col-major; CUBLAS_OP_T ->
|
||||
// (N,K) input[M,K]_row ldb=K -> cuBLAS sees (K,M) col-major; CUBLAS_OP_N
|
||||
// -> (K,M) out[M,N]_row ldc=N -> cuBLAS sees (N,M) col-major (written as
|
||||
// output^T)
|
||||
// cuBLAS: C(N,M) = weight(N,K) @ input(K,M) => C^T = output[M,N]
|
||||
// params: m=N, n=M, k=K, lda=K (weight), ldb=K (input), ldc=N (output)
|
||||
|
||||
torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
|
||||
torch::Tensor const& weight) {
|
||||
TORCH_CHECK(input.dtype() == torch::kBFloat16,
|
||||
"router_gemm_bf16_fp32: input must be bfloat16");
|
||||
TORCH_CHECK(weight.dtype() == torch::kBFloat16,
|
||||
"router_gemm_bf16_fp32: weight must be bfloat16");
|
||||
TORCH_CHECK(input.dim() == 2 && weight.dim() == 2,
|
||||
"router_gemm_bf16_fp32: input and weight must be 2-D");
|
||||
TORCH_CHECK(input.size(1) == weight.size(1),
|
||||
"router_gemm_bf16_fp32: inner dimensions must match");
|
||||
|
||||
int64_t const M = input.size(0);
|
||||
int64_t const N = weight.size(0);
|
||||
int64_t const K = input.size(1);
|
||||
|
||||
auto out = torch::empty({M, N}, input.options().dtype(torch::kFloat32));
|
||||
|
||||
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
|
||||
TORCH_CUDABLAS_CHECK(
|
||||
cublasSetStream(handle, at::cuda::getCurrentCUDAStream()));
|
||||
|
||||
float const alpha = 1.0f;
|
||||
float const beta = 0.0f;
|
||||
|
||||
TORCH_CUDABLAS_CHECK(cublasGemmEx(
|
||||
handle, CUBLAS_OP_T, CUBLAS_OP_N, static_cast<int>(N),
|
||||
static_cast<int>(M), static_cast<int>(K), &alpha, weight.data_ptr(),
|
||||
CUDA_R_16BF, static_cast<int>(K), input.data_ptr(), CUDA_R_16BF,
|
||||
static_cast<int>(K), &beta, out.data_ptr(), CUDA_R_32F,
|
||||
static_cast<int>(N), CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
|
||||
|
||||
return out;
|
||||
}
|
||||
@@ -60,15 +60,6 @@ __device__ __forceinline__ float toFloat(T value) {
|
||||
}
|
||||
}
|
||||
|
||||
#define FINAL_MASK 0xffffffff
|
||||
template <typename T>
|
||||
__inline__ __device__ T warpReduceSum(T val) {
|
||||
#pragma unroll
|
||||
for (int mask = 16; mask > 0; mask >>= 1)
|
||||
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
|
||||
return val;
|
||||
}
|
||||
|
||||
// ====================== TopK softplus_sqrt things
|
||||
// ===============================
|
||||
|
||||
@@ -272,8 +263,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
}
|
||||
}
|
||||
// Compute per-thread scale (using warp reduction when renormalizing).
|
||||
// THREADS_PER_ROW-parameterized butterfly works for both warp sizes (32
|
||||
// on CUDA, 64 on ROCm CDNA) and any THREADS_PER_ROW the dispatch picks.
|
||||
if (renormalize) {
|
||||
selected_sum = warpReduceSum(selected_sum);
|
||||
#pragma unroll
|
||||
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
|
||||
selected_sum +=
|
||||
VLLM_SHFL_XOR_SYNC_WIDTH(selected_sum, mask, THREADS_PER_ROW);
|
||||
}
|
||||
}
|
||||
float scale = static_cast<float>(routed_scaling_factor);
|
||||
if (renormalize) {
|
||||
@@ -544,7 +541,6 @@ void topkGatingSoftplusSqrtKernelLauncher(
|
||||
const IndType* tid2eid, cudaStream_t stream) {
|
||||
static constexpr int WARPS_PER_TB = 4;
|
||||
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
|
||||
#ifndef USE_ROCM
|
||||
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
|
||||
// elements can be loaded by a warp
|
||||
static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
|
||||
@@ -552,6 +548,19 @@ void topkGatingSoftplusSqrtKernelLauncher(
|
||||
std::is_same_v<InputType, __half>)
|
||||
? 4
|
||||
: 8;
|
||||
// Narrower LDG (ELTS_PER_LDG=1) used by 192/320/448/576 on ROCm WARP_SIZE=64
|
||||
// where ELTS_PER_LDG=2 fails the EXPERTS%(ELTS_PER_LDG*WARP_SIZE)==0 check.
|
||||
// On CUDA WARP_SIZE=32 the wider LDG already aligns, so the alias collapses
|
||||
// back to BYTES_PER_LDG_MULTIPLE_64 — no behavioral change for CUDA.
|
||||
#ifdef USE_ROCM
|
||||
static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW =
|
||||
(std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>)
|
||||
? 2
|
||||
: 4;
|
||||
#else
|
||||
static constexpr int BYTES_PER_LDG_MULTIPLE_64_NARROW =
|
||||
BYTES_PER_LDG_MULTIPLE_64;
|
||||
#endif
|
||||
switch (num_experts) {
|
||||
case 1:
|
||||
@@ -584,27 +593,29 @@ void topkGatingSoftplusSqrtKernelLauncher(
|
||||
case 512:
|
||||
LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
// (CUDA only) support multiples of 64 when num_experts is not power of 2.
|
||||
// ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of
|
||||
// num_experts, alternatively we can test 4 bytes loading and enable it in
|
||||
// future.
|
||||
#ifndef USE_ROCM
|
||||
// Multiples of 64 that are not powers of 2. The kernel requires
|
||||
// EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0. With ELTS_PER_LDG=2
|
||||
// (BYTES_PER_LDG_MULTIPLE_64), this holds for all five values on CUDA
|
||||
// WARP_SIZE=32 but only for 384 on ROCm WARP_SIZE=64. The other four
|
||||
// use BYTES_PER_LDG_MULTIPLE_64_NARROW (ELTS_PER_LDG=1), which
|
||||
// satisfies the assertion for any multiple of 64 on either backend;
|
||||
// on CUDA the narrow alias collapses back to the wider load, so CUDA
|
||||
// behavior is unchanged.
|
||||
case 192:
|
||||
LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
|
||||
break;
|
||||
case 320:
|
||||
LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
|
||||
break;
|
||||
case 384:
|
||||
LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
break;
|
||||
case 448:
|
||||
LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
|
||||
break;
|
||||
case 576:
|
||||
LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64_NARROW);
|
||||
break;
|
||||
#endif
|
||||
default: {
|
||||
TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
|
||||
}
|
||||
|
||||
@@ -16,14 +16,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
|
||||
"bias) -> ()");
|
||||
m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
m.def(
|
||||
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, float "
|
||||
"routed_scaling_factor, Tensor? "
|
||||
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
|
||||
m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
|
||||
#endif
|
||||
|
||||
// Calculate the result of moe by summing up the partial results
|
||||
// from all selected experts.
|
||||
m.def("moe_sum(Tensor input, Tensor! output) -> ()");
|
||||
@@ -133,10 +132,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
|
||||
"Tensor)");
|
||||
m.impl("grouped_topk", torch::kCUDA, &grouped_topk);
|
||||
|
||||
// cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16)
|
||||
m.def("router_gemm_bf16_fp32(Tensor input, Tensor weight) -> Tensor");
|
||||
m.impl("router_gemm_bf16_fp32", torch::kCUDA, &router_gemm_bf16_fp32);
|
||||
|
||||
// DeepSeek V3 optimized router GEMM for SM90+
|
||||
m.def("dsv3_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
|
||||
// conditionally compiled so impl registration is in source file
|
||||
|
||||
@@ -887,27 +887,14 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
|
||||
uint32_t* shared_ordered =
|
||||
reinterpret_cast<uint32_t*>(smem_raw + kFixedSmemLarge);
|
||||
|
||||
// RadixRowState for multi-CTA cooperative radix
|
||||
// RadixRowState for multi-CTA cooperative radix.
|
||||
// Zero-initialization is done host-side via cudaMemsetAsync in topk.cu
|
||||
// before launch — that gives a stream-ordered happens-before edge for all
|
||||
// CTAs, which the previous in-kernel init (CTA-0 only + intra-CTA
|
||||
// __syncthreads) did not provide and which manifested as a race against
|
||||
// CTA-1+'s first red_release on arrival_counter.
|
||||
RadixRowState* state = ¶ms.row_states[group_id];
|
||||
|
||||
// -- Initialize RadixRowState (only needed if large rows exist) --
|
||||
if (params.max_seq_len > RADIX_THRESHOLD) {
|
||||
if (cta_in_group == 0) {
|
||||
for (uint32_t buf = 0; buf < 3; buf++) {
|
||||
for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) {
|
||||
state->histogram[buf][i] = 0;
|
||||
}
|
||||
}
|
||||
if (tx == 0) {
|
||||
state->remaining_k = 0;
|
||||
state->prefix = 0;
|
||||
state->arrival_counter = 0;
|
||||
state->output_counter = 0;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
int barrier_phase = 0;
|
||||
const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups;
|
||||
|
||||
|
||||
@@ -153,6 +153,29 @@ void launch_persistent_topk(const torch::Tensor& logits,
|
||||
TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
|
||||
"workspace too small, need ", state_bytes, " bytes");
|
||||
|
||||
// Zero the per-group RadixRowState region before launch — only when the
|
||||
// radix path will actually run (max_seq_len > RADIX_THRESHOLD). The
|
||||
// RadixRowState fields (arrival_counter, histograms) are only touched by
|
||||
// radix_topk; the decode/medium paths inside the persistent kernel
|
||||
// operate purely in shared memory and never read these globals, so a
|
||||
// stale workspace is harmless for them.
|
||||
//
|
||||
// Why we need the memset (when needs_cooperative is true):
|
||||
// 1. arrival_counter accumulates within a launch and is never reset,
|
||||
// so a prior call leaves it at a large positive value. Without this
|
||||
// reset, the very first wait_ge in the next call sees counter >>
|
||||
// target and returns instantly, breaking the barrier.
|
||||
// 2. The previous in-kernel init only ran in CTA-0 with intra-CTA
|
||||
// __syncthreads(), so it had no happens-before edge to CTA-1+'s
|
||||
// first red_release. cudaMemsetAsync is stream-ordered: the zero
|
||||
// is globally visible before any CTA runs.
|
||||
if (needs_cooperative) {
|
||||
cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr<uint8_t>(), 0,
|
||||
state_bytes, stream);
|
||||
TORCH_CHECK(mz_err == cudaSuccess,
|
||||
"row_states memset failed: ", cudaGetErrorString(mz_err));
|
||||
}
|
||||
|
||||
P::PersistentTopKParams params;
|
||||
params.input = logits.data_ptr<float>();
|
||||
params.output = output.data_ptr<int32_t>();
|
||||
|
||||
@@ -183,7 +183,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"int forced_token_heads_per_warp=-1) -> ()");
|
||||
ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
|
||||
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
|
||||
// kernel launch.
|
||||
@@ -194,7 +193,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"float eps, int cache_block_size) -> ()");
|
||||
ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
|
||||
&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
|
||||
#endif
|
||||
|
||||
// Apply repetition penalties to logits in-place
|
||||
ops.def(
|
||||
|
||||
+67
-26
@@ -41,6 +41,13 @@ ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
|
||||
# Using cuda base image with minimal dependencies necessary for JIT compilation (FlashInfer, DeepGEMM, EP kernels)
|
||||
ARG FINAL_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-base-ubuntu${UBUNTU_VERSION}
|
||||
|
||||
# OS family of BUILD_BASE_IMAGE. Controls package manager (apt vs dnf) and
|
||||
# Python bootstrap. Set to "manylinux" alongside a manylinux build base such
|
||||
# as pytorch/manylinux2_28-builder:cuda13.0 to produce wheels with a glibc
|
||||
# 2.28 floor (matches PyTorch's own published wheels). Default stays on
|
||||
# Ubuntu for backwards compatibility.
|
||||
ARG BUILD_OS=ubuntu
|
||||
|
||||
# By parameterizing the Deadsnakes repository URL, we allow third-party to use
|
||||
# their own mirror. When doing so, we don't benefit from the transparent
|
||||
# installation of the GPG key of the PPA, as done by add-apt-repository, so we
|
||||
@@ -94,35 +101,64 @@ FROM ${BUILD_BASE_IMAGE} AS base
|
||||
|
||||
ARG CUDA_VERSION
|
||||
ARG PYTHON_VERSION
|
||||
ARG BUILD_OS
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
# Install system dependencies including build tools
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ccache \
|
||||
software-properties-common \
|
||||
git \
|
||||
curl \
|
||||
sudo \
|
||||
python3-pip \
|
||||
libibverbs-dev \
|
||||
# Upgrade to GCC 10 to avoid https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92519
|
||||
# as it was causing spam when compiling the CUTLASS kernels
|
||||
gcc-10 \
|
||||
g++-10 \
|
||||
&& update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 110 --slave /usr/bin/g++ g++ /usr/bin/g++-10 \
|
||||
# Install python dev headers if available (needed for cmake FindPython on Ubuntu 24.04
|
||||
# which ships cmake 3.28 and requires Development.SABIModule; silently skipped on
|
||||
# Ubuntu 20.04/22.04 where python3.x-dev is not available without a PPA)
|
||||
&& (apt-get install -y --no-install-recommends python${PYTHON_VERSION}-dev 2>/dev/null || true) \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& curl -LsSf https://astral.sh/uv/install.sh | sh \
|
||||
&& $HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION} \
|
||||
# Install system dependencies including build tools.
|
||||
# The Ubuntu path uses apt + deadsnakes-via-uv for Python; the manylinux path
|
||||
# (AlmaLinux 8, e.g. pytorch/manylinux2_28-builder) uses dnf and the Python
|
||||
# interpreters pre-installed at /opt/python/cpXY-cpXY/.
|
||||
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
|
||||
# rdma-core-devel provides libibverbs headers; ccache lives in EPEL,
|
||||
# which the pytorch manylinux image already enables. git/curl/sudo
|
||||
# are typically pre-installed but listed defensively.
|
||||
dnf install -y --setopt=install_weak_deps=False \
|
||||
ccache \
|
||||
git \
|
||||
curl \
|
||||
sudo \
|
||||
rdma-core-devel \
|
||||
&& dnf clean all \
|
||||
&& rm -rf /var/cache/dnf; \
|
||||
else \
|
||||
apt-get update -y \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
ccache \
|
||||
software-properties-common \
|
||||
git \
|
||||
curl \
|
||||
sudo \
|
||||
python3-pip \
|
||||
libibverbs-dev \
|
||||
# Upgrade to GCC 10 to avoid https://gcc.gnu.org/bugzilla/show_bug.cgi?id=92519
|
||||
# as it was causing spam when compiling the CUTLASS kernels
|
||||
gcc-10 \
|
||||
g++-10 \
|
||||
&& update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-10 110 --slave /usr/bin/g++ g++ /usr/bin/g++-10 \
|
||||
# Install python dev headers if available (needed for cmake FindPython on Ubuntu 24.04
|
||||
# which ships cmake 3.28 and requires Development.SABIModule; silently skipped on
|
||||
# Ubuntu 20.04/22.04 where python3.x-dev is not available without a PPA)
|
||||
&& (apt-get install -y --no-install-recommends python${PYTHON_VERSION}-dev 2>/dev/null || true) \
|
||||
&& rm -rf /var/lib/apt/lists/*; \
|
||||
fi
|
||||
|
||||
# Install uv and bootstrap /opt/venv. Both paths converge on /opt/venv so all
|
||||
# downstream stages stay distro-agnostic.
|
||||
RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
|
||||
&& if [ "${BUILD_OS}" = "manylinux" ]; then \
|
||||
# manylinux images ship Python at /opt/python/cpXY-cpXY/; point uv
|
||||
# at the matching interpreter rather than letting it fetch one.
|
||||
PYV_NODOT=$(echo ${PYTHON_VERSION} | tr -d '.') \
|
||||
&& MANYLINUX_PY=/opt/python/cp${PYV_NODOT}-cp${PYV_NODOT}/bin/python${PYTHON_VERSION} \
|
||||
&& $HOME/.local/bin/uv venv /opt/venv --python "$MANYLINUX_PY"; \
|
||||
else \
|
||||
$HOME/.local/bin/uv venv /opt/venv --python ${PYTHON_VERSION}; \
|
||||
fi \
|
||||
&& rm -f /usr/bin/python3 /usr/bin/python3-config /usr/bin/pip \
|
||||
&& ln -s /opt/venv/bin/python3 /usr/bin/python3 \
|
||||
&& ln -s /opt/venv/bin/python3-config /usr/bin/python3-config \
|
||||
&& ln -s /opt/venv/bin/pip /usr/bin/pip \
|
||||
&& ln -sf /opt/venv/bin/python3 /usr/bin/python3 \
|
||||
&& ln -sf /opt/venv/bin/python3-config /usr/bin/python3-config \
|
||||
&& ln -sf /opt/venv/bin/pip /usr/bin/pip \
|
||||
&& python3 --version && python3 -m pip --version
|
||||
|
||||
# Activate virtual environment and add uv to PATH
|
||||
@@ -433,6 +469,7 @@ FROM base AS dev
|
||||
ARG PIP_INDEX_URL UV_INDEX_URL
|
||||
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
|
||||
ARG PYTORCH_CUDA_INDEX_BASE_URL
|
||||
ARG BUILD_OS
|
||||
|
||||
# This timeout (in seconds) is necessary when installing some dependencies via uv since it's likely to time out
|
||||
# Reference: https://github.com/astral-sh/uv/pull/1694
|
||||
@@ -442,7 +479,11 @@ ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE=copy
|
||||
|
||||
# Install libnuma-dev, required by fastsafetensors (fixes #20384)
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends libnuma-dev && rm -rf /var/lib/apt/lists/*
|
||||
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
|
||||
dnf install -y numactl-devel && dnf clean all && rm -rf /var/cache/dnf; \
|
||||
else \
|
||||
apt-get update && apt-get install -y --no-install-recommends libnuma-dev && rm -rf /var/lib/apt/lists/*; \
|
||||
fi
|
||||
|
||||
|
||||
# We can specify the standard or nightly build of PyTorch
|
||||
|
||||
@@ -124,9 +124,9 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
|
||||
|
||||
# RIXL/UCX build stages
|
||||
FROM base AS build_rixl
|
||||
ARG RIXL_BRANCH="bf4a7214"
|
||||
ARG RIXL_BRANCH="39be1de8"
|
||||
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
|
||||
ARG UCX_BRANCH="7009d7a1"
|
||||
ARG UCX_BRANCH="bfb51733"
|
||||
ARG UCX_REPO="https://github.com/openucx/ucx.git"
|
||||
ENV ROCM_PATH=/opt/rocm
|
||||
ENV UCX_HOME=/usr/local/ucx
|
||||
@@ -192,6 +192,7 @@ RUN cd /opt/rixl && \
|
||||
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
|
||||
contrib/build-wheel.sh && \
|
||||
mkdir -p /app/install && \
|
||||
_ucx_install_dir=${UCX_HOME} \
|
||||
./contrib/build-wheel.sh \
|
||||
--output-dir /app/install \
|
||||
--rocm-dir ${ROCM_PATH} \
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.1-complete
|
||||
ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.2-complete
|
||||
ARG TRITON_BRANCH="ba5c1517"
|
||||
ARG TRITON_REPO="https://github.com/ROCm/triton.git"
|
||||
ARG PYTORCH_BRANCH="8514f051" # release/2.10 as of 3/17
|
||||
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
|
||||
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
|
||||
ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="v0.1.10.post3"
|
||||
ARG AITER_BRANCH="v0.1.12.post2"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="v1.1.0"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
@@ -104,6 +104,28 @@ ENV SCCACHE_REGION=${USE_SCCACHE:+${SCCACHE_REGION_NAME}}
|
||||
ENV SCCACHE_S3_NO_CREDENTIALS=${USE_SCCACHE:+${SCCACHE_S3_NO_CREDENTIALS}}
|
||||
ENV SCCACHE_IDLE_TIMEOUT=${USE_SCCACHE:+0}
|
||||
|
||||
# torch profiler hotfix for 7.2.2: rebuild CLR with https://github.com/ROCm/rocm-systems/pull/5062
|
||||
# will be removed once we move to ROCm 7.2.3
|
||||
RUN apt-get update && apt-get install -y rocm-llvm-dev
|
||||
RUN pip install CppHeaderParser
|
||||
RUN git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems /tmp/rocm-systems \
|
||||
&& cd /tmp/rocm-systems \
|
||||
&& git sparse-checkout init --cone \
|
||||
&& git sparse-checkout set projects/hip projects/clr \
|
||||
&& git checkout 35e8c7bf8911862e5389509800e65fdf125412b3 \
|
||||
&& export CLR_DIR=/tmp/rocm-systems/projects/clr \
|
||||
&& export HIP_DIR=/tmp/rocm-systems/projects/hip \
|
||||
&& mkdir -p $CLR_DIR/build && cd $CLR_DIR/build \
|
||||
&& cmake \
|
||||
-DHIP_COMMON_DIR=$HIP_DIR \
|
||||
-DCMAKE_PREFIX_PATH="/opt/rocm/" \
|
||||
-DCLR_BUILD_HIP=ON \
|
||||
-DCLR_BUILD_OCL=OFF \
|
||||
-DHIP_PLATFORM=amd \
|
||||
.. \
|
||||
&& make -j$(nproc) \
|
||||
&& make install \
|
||||
&& rm -rf /tmp/rocm-systems
|
||||
|
||||
###
|
||||
### Triton Build
|
||||
@@ -153,8 +175,6 @@ RUN git clone ${PYTORCH_REPO} pytorch
|
||||
RUN cd pytorch && git checkout ${PYTORCH_BRANCH}
|
||||
RUN cd pytorch \
|
||||
&& pip install -r requirements.txt && git submodule update --init --recursive
|
||||
RUN cd pytorch/third_party/kineto \
|
||||
&& git remote add rocm https://github.com/ROCm/kineto && git fetch rocm && git checkout 2d73be3
|
||||
RUN cd pytorch && python3 tools/amd_build/build_amd.py \
|
||||
&& if [ "$USE_SCCACHE" = "1" ]; then \
|
||||
export HIP_CLANG_PATH=/opt/sccache-wrappers \
|
||||
|
||||
@@ -16,6 +16,9 @@
|
||||
"FINAL_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
|
||||
},
|
||||
"BUILD_OS": {
|
||||
"default": "ubuntu"
|
||||
},
|
||||
"GET_PIP_URL": {
|
||||
"default": "https://bootstrap.pypa.io/get-pip.py"
|
||||
},
|
||||
|
||||
@@ -60,9 +60,19 @@ the failure?
|
||||
|
||||
## Logs Wrangling
|
||||
|
||||
Download the full log file from Buildkite locally.
|
||||
Download a job's log (no Buildkite login required):
|
||||
|
||||
Strip timestamps and colorization:
|
||||
[.buildkite/scripts/ci-fetch-log.sh](../../../.buildkite/scripts/ci-fetch-log.sh)
|
||||
|
||||
```bash
|
||||
# Find the failing job. Each row's URL is .../builds/<N>#<job_uuid>:
|
||||
gh pr checks <PR> --repo vllm-project/vllm
|
||||
|
||||
# Download + strip timestamps/ANSI in one step:
|
||||
.buildkite/scripts/ci-fetch-log.sh "https://buildkite.com/vllm/ci/builds/<N>#<job_uuid>"
|
||||
```
|
||||
|
||||
To clean an already-downloaded log:
|
||||
|
||||
[.buildkite/scripts/ci-clean-log.sh](../../../.buildkite/scripts/ci-clean-log.sh)
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
[Anyscale](https://www.anyscale.com) is a managed, multi-cloud platform developed by the creators of Ray.
|
||||
|
||||
Anyscale automates the entire lifecycle of Ray clusters in your AWS, GCP, or Azure account, delivering the flexibility of open-source Ray
|
||||
without the operational overhead of maintaining Kubernetes control planes, configuring autoscalers, managing observability stacks, or manually managing head and worker nodes with helper scripts like [examples/online_serving/run_cluster.sh](../../../examples/online_serving/run_cluster.sh).
|
||||
without the operational overhead of maintaining Kubernetes control planes, configuring autoscalers, managing observability stacks, or manually managing head and worker nodes with helper scripts like [examples/ray_serving/run_cluster.sh](../../../examples/ray_serving/run_cluster.sh).
|
||||
|
||||
When serving large language models with vLLM, Anyscale can rapidly provision [production-ready HTTPS endpoints](https://docs.anyscale.com/examples/deploy-ray-serve-llms) or [fault-tolerant batch inference jobs](https://docs.anyscale.com/examples/ray-data-llm).
|
||||
|
||||
|
||||
@@ -155,6 +155,7 @@ Priority is **1 = highest** (tried first).
|
||||
| **Block Sizes** | Supported KV cache block sizes (%N means multiples of N) |
|
||||
| **Head Sizes** | Supported attention head sizes |
|
||||
| **Sink** | Attention sink support (for StreamingLLM) |
|
||||
| **Non-Causal** | Non-causal (bidirectional) attention support for decoder models |
|
||||
| **Sparse** | Sparse attention support (MLA only) |
|
||||
| **MM Prefix** | Multimodal prefix full attention support |
|
||||
| **DCP** | Decode Context Parallelism support (`--decode-context-parallel-size`) |
|
||||
@@ -165,22 +166,22 @@ Priority is **1 = highest** (tried first).
|
||||
|
||||
## Standard Attention (MHA, MQA, GQA) Backends
|
||||
|
||||
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
|
||||
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
|
||||
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A |
|
||||
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
|
||||
| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
|
||||
| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | --------- | --- | --------------- | ------------ |
|
||||
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | ❌ | All | N/A |
|
||||
| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ❌ | ✅ | Decoder | 7.x-9.x |
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
|
||||
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
|
||||
| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | Any |
|
||||
| `TURBOQUANT` | | fp16, bf16 | `turboquant_k8v4`, `turboquant_4bit_nc`, `turboquant_k3v4_nc`, `turboquant_3bit_nc` | 16, 32, 64, 128 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
|
||||
> **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`.
|
||||
>
|
||||
@@ -192,31 +193,35 @@ MLA uses separate backends for prefill and decode phases.
|
||||
|
||||
### Prefill Backends
|
||||
|
||||
The prefill backend is selected at runtime based on hardware and
|
||||
configuration.
|
||||
To explicitly select a prefill backend, use
|
||||
`-ac.mla_prefill_backend=<BACKEND>` (e.g., `FLASH_ATTN`, `FLASHINFER`).
|
||||
Otherwise, the prefill backend is selected automatically at runtime based on
|
||||
hardware and configuration.
|
||||
|
||||
| Backend | Description | Compute Cap. | Enable | Disable | Notes |
|
||||
| ------- | ----------- | ------------ | ------ | ------- | ----- |
|
||||
| TRT-LLM Ragged‡ | TensorRT-LLM ragged attention | 10.x | Default on SM100 | `-ac.use_trtllm_ragged_deepseek_prefill=0` | DeepSeek R1 dims only |
|
||||
| FlashInfer | FlashInfer CUTLASS backend | 10.x | `-ac.disable_flashinfer_prefill=0` | `-ac.disable_flashinfer_prefill=1` | DeepSeek R1 dims only |
|
||||
| cuDNN | cuDNN-based attention | 10.x | `-ac.use_cudnn_prefill=1` | `-ac.use_cudnn_prefill=0` | |
|
||||
| FlashAttention | FlashAttention varlen (FA2/FA3) | Any | Default fallback | Use other backends | FA3 on SM90, FA2 otherwise |
|
||||
| Backend | Description | Dtypes | Compute Cap. | Notes |
|
||||
| ------- | ----------- | ------ | ------------ | ----- |
|
||||
| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | FA4 on SM100+, FA3 on SM90, FA2 otherwise |
|
||||
| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | DeepSeek R1 dims only |
|
||||
| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | DeepSeek R1 dims only |
|
||||
|
||||
> **‡** TRT-LLM Ragged is the default on Blackwell (SM100).
|
||||
> On other GPUs, FlashAttention is used as the default.
|
||||
|
||||
### Decode Backends
|
||||
|
||||
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
|
||||
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
|
||||
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
|
||||
MLA decode backends are selected using the standard
|
||||
`-ac.backend=<BACKEND>` argument (e.g., `FLASHMLA`, `TRITON_MLA`).
|
||||
|
||||
| Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Non-Causal | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------ | --------- | ----------- | ---------- | ---- | ---------- | ------ | --------- | --- | --------------- | ------------ |
|
||||
| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
|
||||
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 1, 64 | Any | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ❌ | ✅ | ❌ | ❌ | Decoder | Any |
|
||||
|
||||
@@ -86,9 +86,11 @@ 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` | ✅︎ | ✅︎ |
|
||||
|
||||
!!! note
|
||||
Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
|
||||
For Qwen2.5-VL only FA2 and FA3 has been tested.
|
||||
|
||||
## Configuration
|
||||
|
||||
|
||||
@@ -5,12 +5,14 @@ TL;DR:
|
||||
- use tlparse to acquire torch.compile logs. Include these logs in bug reports and/or support asks.
|
||||
- The vLLM-torch.compile integration is multiple pieces. vLLM exposes flags to turn off each piece:
|
||||
|
||||
| Online Flag | Offline Flag | Result |
|
||||
| ----------- | ------------ | ------ |
|
||||
| --enforce-eager | enforce_eager=True | Turn off torch.compile and CUDAGraphs |
|
||||
| -cc.mode=0 | mode=CompilationMode.NONE | Turn off torch.compile only |
|
||||
| -cc.cudagraph_mode=NONE | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE) | Turn off CUDAGraphs only |
|
||||
| -cc.backend=eager | compilation_config=CompilationConfig(backend='eager') | Turn off TorchInductor |
|
||||
| Online Flag | Offline Flag | Result |
|
||||
|--------------------------------|--------------------------------------------------------------------------------|------------------------------------------------------|
|
||||
| --enforce-eager | enforce_eager=True | Turn off torch.compile and CUDAGraphs |
|
||||
| -cc.mode=0 | compilation_config=CompilationConfig(mode=CompilationMode.NONE) | Turn off torch.compile only |
|
||||
| -cc.mode=1 | compilation_config=CompilationConfig(mode=CompilationMode.STOCK_TORCH_COMPILE) | Turn off vLLM-compile modifications to torch.compile |
|
||||
| -cc.cudagraph_mode=NONE | compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE) | Turn off CUDAGraphs only |
|
||||
| -cc.backend=eager | compilation_config=CompilationConfig(backend='eager') | Turn off TorchInductor |
|
||||
| -cc.ir_enable_torch_wrap=False | compilation_config=CompilationConfig(ir_enable_torch_wrap=False) | Turn off vLLM IR wrapping |
|
||||
|
||||
## vLLM-torch.compile overview
|
||||
|
||||
@@ -22,7 +24,7 @@ Most notably, vLLM-compile is NOT torch.compile, it is a custom compiler built u
|
||||
|
||||
- Given a model, we do a full graph capture via TorchDynamo that is dynamic on the batch size (number of tokens)
|
||||
- vLLM then optionally splits and/or specializes this graph and then uses TorchInductor to compile each graph into a compiled artifact.
|
||||
This step may use vLLM custom Inductor passes to further optimize the graph.
|
||||
This step may use vLLM custom Inductor passes to further optimize the graph. This includes vLLM IR lowering to remove dispatch overhead.
|
||||
- The compiled artifact is saved to vLLM's compile cache so that it can be loaded in the future.
|
||||
- vLLM applies CUDAGraphs to reduce CPU overheads.
|
||||
|
||||
@@ -34,6 +36,7 @@ For more details on the design, please see the following resources:
|
||||
|
||||
- [Introduction to vLLM-torch.compile blogpost](https://blog.vllm.ai/2025/08/20/torch-compile.html)
|
||||
- [vLLM-torch.compile integration design](./torch_compile.md)
|
||||
- [vLLM IR design](./vllm_ir.md)
|
||||
- [vLLM Office Hours #26](https://www.youtube.com/live/xLyxc7hxCJc?si=Xulo9pe53C6ywf0V&t=561)
|
||||
- [Talk at PyTorch Conference 2025](https://youtu.be/1wV1ESbGrVQ?si=s1GqymUfwiwOrDTg&t=725)
|
||||
|
||||
@@ -117,6 +120,21 @@ from vllm.config.compilation import CompilationConfig, CUDAGraphMode
|
||||
LLM(model, compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE))
|
||||
```
|
||||
|
||||
vLLM IR makes heavy use of the compilation pipeline, from functionalization, custom fusions, and lowering.
|
||||
To turn that off and capture eager-mode dispatching behavior of vLLM IR, run with `ir_enable_torch_wrap=False`.
|
||||
IR torch wrap is only enabled by default when using `mode=VLLM_COMPILE` and `backend="inductor"` (default).
|
||||
|
||||
```sh
|
||||
# Online
|
||||
vllm serve -cc.ir_enable_torch_wrap=False
|
||||
```
|
||||
|
||||
```py
|
||||
# Offline
|
||||
from vllm.config.compilation import CompilationConfig
|
||||
LLM(model, compilation_config=CompilationConfig(ir_enable_torch_wrap=False))
|
||||
```
|
||||
|
||||
## Debugging TorchDynamo
|
||||
|
||||
vLLM requires model code be capturable into a full graph via TorchDynamo (torch.compile's frontend).
|
||||
|
||||
@@ -36,7 +36,7 @@ th {
|
||||
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
|
||||
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
|
||||
!!! info "Table key"
|
||||
1. All types: mxfp4, nvfp4, int4, int8, fp8
|
||||
|
||||
@@ -88,7 +88,7 @@ pip install "vllm>=0.9.2"
|
||||
#### Proxy (e.g. 10.0.1.1)
|
||||
|
||||
```shell
|
||||
cd {your vllm directory}/examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/
|
||||
cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
|
||||
python3 disagg_proxy_p2p_nccl_xpyd.py &
|
||||
```
|
||||
|
||||
@@ -181,7 +181,7 @@ python3 disagg_proxy_p2p_nccl_xpyd.py &
|
||||
#### Proxy (e.g. 10.0.1.1)
|
||||
|
||||
```shell
|
||||
cd {your vllm directory}/examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/
|
||||
cd {your vllm directory}/examples/disaggregated/p2p_nccl_xpyd/
|
||||
python3 disagg_proxy_p2p_nccl_xpyd.py &
|
||||
```
|
||||
|
||||
|
||||
@@ -0,0 +1,615 @@
|
||||
# vLLM IR: Functional Intermediate Representation
|
||||
|
||||
## Motivation
|
||||
|
||||
vLLM IR is a **functional intermediate representation (IR)** that fills the gap between
|
||||
low-level `torch` ops and vLLM layers like `RMSNorm` and quantization operators,
|
||||
By separating operator **semantics** from the **implementation** and **dispatching**,
|
||||
vLLM IR simplifies both compilation and kernel registration & dispatching simultaneously.
|
||||
It operates as a **dialect** in the torch FX representation, allowing full interoperability
|
||||
with “regular” torch ops & custom torch ops/kernels, as well as a piecewise migration from
|
||||
the previous `CustomOp` approach.
|
||||
|
||||
Key design principles:
|
||||
|
||||
- **Eager-compile consistency**: identical behavior (barring minor numerics) in eager and compiled modes
|
||||
- **Simple, transparent, yet powerful kernel selection**: good visibility and control allowing easy debugging
|
||||
- **Convention over configuration**: near-zero boilerplate required to register ops and implementations
|
||||
- **Extensibility**: ops and implementations can be registered anywhere, in-tree or out-of-tree
|
||||
- **Interoperability**: fully compatible with “regular” torch ops & custom torch ops/kernels,
|
||||
reducing developer friction and allowing piecewise migration
|
||||
|
||||
The clean semantics/implementation separation enables a unified and extensible dispatching mechanism,
|
||||
allowing multiple kernels per-platform and powerful kernel selection. The separation also facilitates
|
||||
cleaner testing and benchmarking, removing much of the boilerplate standard for legacy approaches.
|
||||
|
||||
By delaying kernel selection until late in the compilation process, the compiler can operate on
|
||||
a higher-level representation, which has the following main benefits:
|
||||
|
||||
- Pattern matching in fusion/transformation passes only requires a single, simple pattern per op
|
||||
- OOT compiler backends can lower from the higher-level representation (in-progress)
|
||||
- The compiler can autotune over available implementations (future feature)
|
||||
|
||||
## Quick Overview
|
||||
|
||||
### Declaring an IR Operation
|
||||
|
||||
IR operations are declared using the `@register_op` decorator with a native PyTorch implementation that defines the op's semantics:
|
||||
|
||||
```python
|
||||
# vllm/ir/ops/layernorm.py
|
||||
from torch import Tensor
|
||||
from vllm.ir import register_op
|
||||
|
||||
@register_op
|
||||
def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor:
|
||||
"""Weighted root-mean-square layer normalization"""
|
||||
orig_dtype = x.dtype
|
||||
x = x.to(torch.float32)
|
||||
x_var = x if variance_size is None else x[..., :variance_size]
|
||||
variance = x_var.pow(2).mean(dim=-1, keepdim=True)
|
||||
x = x * torch.rsqrt(variance + epsilon)
|
||||
x = x.to(orig_dtype)
|
||||
if weight is not None:
|
||||
x = x * weight
|
||||
return x
|
||||
```
|
||||
|
||||
The native implementation serves three purposes:
|
||||
|
||||
1. **Semantic definition**: Specifies the exact semantics of the operation, including shapes and strides
|
||||
2. **Default implementation**: Used when no other (better) implementation is available
|
||||
3. **Reference for testing**: Other implementations must match these semantics
|
||||
|
||||
### Registering Implementations
|
||||
|
||||
Kernel implementations are registered using the `register_impl` decorator on the IR op object:
|
||||
|
||||
```python
|
||||
# vllm/kernels/vllm_c.py
|
||||
from vllm import ir
|
||||
|
||||
rms_norm_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None
|
||||
|
||||
@ir.ops.rms_norm.register_impl("vllm_c", supports_args=rms_norm_no_var, supported=current_platform.is_cuda_alike())
|
||||
def rms_norm(x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None) -> Tensor:
|
||||
output = torch.empty_like(x)
|
||||
torch.ops._C.rms_norm(output, x, weight, epsilon)
|
||||
return output
|
||||
```
|
||||
|
||||
Implementations can specify:
|
||||
|
||||
- `supported`: Static boolean indicating if this implementation is available
|
||||
- `supports_args`: Function checking if the implementation supports specific arguments
|
||||
- `inplace`: Whether this implementation reuses input memory for outputs
|
||||
|
||||
### Using IR Operations in Models
|
||||
|
||||
IR operations are imported and called directly in model code:
|
||||
|
||||
```python
|
||||
# vllm/model_executor/layers/layernorm.py
|
||||
from vllm import ir
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, hidden_size: int, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
self.variance_epsilon = eps
|
||||
|
||||
def forward(self, x: Tensor, residual: Tensor | None = None):
|
||||
if residual is None:
|
||||
return ir.ops.rms_norm(x, self.weight, self.variance_epsilon)
|
||||
|
||||
# Use maybe_inplace overload to allow implementation to reuse input memory for outputs
|
||||
# (using x or residual after this call is undefined behavior)
|
||||
return ir.ops.fused_add_rms_norm.maybe_inplace(
|
||||
x, residual, self.weight, self.variance_epsilon
|
||||
)
|
||||
```
|
||||
|
||||
### Configuring Kernel Selection
|
||||
|
||||
Kernel selection is controlled via priority lists in the configuration.
|
||||
Priority lists specify the order in which implementations are considered,
|
||||
with the first supported implementation being selected.
|
||||
This includes the static support check (`supported=...`) and
|
||||
the dynamic arg support check (`supports_args=...`).
|
||||
|
||||
#### Command Line Configuration
|
||||
|
||||
Use `--ir-op-priority.<op_name>=<provider1>,<provider2>,...`:
|
||||
|
||||
```bash
|
||||
# CUDA: Use vllm_c implementation for rms_norm
|
||||
vllm serve meta-llama/Llama-3.2-1B \
|
||||
--ir-op-priority.rms_norm=vllm_c
|
||||
|
||||
# ROCm: Try aiter first, fall back to vllm_c, then native
|
||||
vllm serve meta-llama/Llama-3.2-1B \
|
||||
--ir-op-priority.rms_norm=aiter,vllm_c,native
|
||||
|
||||
# Configure multiple operations
|
||||
vllm serve meta-llama/Llama-3.2-1B \
|
||||
--ir-op-priority.rms_norm=vllm_c \
|
||||
--ir-op-priority.fused_add_rms_norm=vllm_c
|
||||
```
|
||||
|
||||
#### Python Configuration
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from vllm.config import VllmConfig, KernelConfig
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Llama-3.2-1B",
|
||||
vllm_config=VllmConfig(
|
||||
kernel_config=KernelConfig(
|
||||
ir_op_priority={
|
||||
"rms_norm": ["vllm_c", "native"],
|
||||
"fused_add_rms_norm": ["vllm_c", "native"],
|
||||
}
|
||||
)
|
||||
)
|
||||
)
|
||||
```
|
||||
|
||||
#### Platform Defaults
|
||||
|
||||
Each platform provides default priority lists that are automatically applied:
|
||||
|
||||
```python
|
||||
# CUDA/XPU/ROCm platform defaults (when compiling with Inductor)
|
||||
{
|
||||
"rms_norm": ["native"], # Native torch is default
|
||||
"fused_add_rms_norm": ["native"],
|
||||
}
|
||||
|
||||
# CUDA platform defaults (eager or Dynamo-only)
|
||||
{
|
||||
"rms_norm": ["vllm_c", "native"],
|
||||
"fused_add_rms_norm": ["vllm_c", "native"],
|
||||
}
|
||||
|
||||
# ROCm platform defaults (future - currently same as CUDA)
|
||||
{
|
||||
"rms_norm": ["aiter", "vllm_c", "native"],
|
||||
"fused_add_rms_norm": ["aiter", "vllm_c", "native"],
|
||||
}
|
||||
|
||||
# XPU platform defaults (eager or Dynamo-only)
|
||||
{
|
||||
"rms_norm": ["xpu_kernels", "native"],
|
||||
"fused_add_rms_norm": ["xpu_kernels", "native"],
|
||||
}
|
||||
```
|
||||
|
||||
User-specified priorities are prepended to platform defaults,
|
||||
so you only need to specify the out-of-order implementations,
|
||||
other implementations are appended automatically.
|
||||
|
||||
## Compilation Pipeline
|
||||
|
||||
vLLM IR heavily customizes the `torch.compile`-based compilation process to allow custom compile
|
||||
passes to operate on high-level IR while still producing efficient low-level code at the end.
|
||||
The compilation pipeline consists of several stages:
|
||||
|
||||
### 1. Dynamo Tracing
|
||||
|
||||
When `torch.compile` traces the model's forward pass, vLLM IR operations appear as custom operations
|
||||
in the `vllm_ir` torch library. These operations are opaque to Dynamo, meaning they appear directly
|
||||
in the FX graph without decomposition:
|
||||
|
||||
```python
|
||||
# Python code (epsilon=1e-5)
|
||||
x1 = ir.ops.rms_norm(x, weight, epsilon)
|
||||
x2, residual_out = ir.ops.fused_add_rms_norm.maybe_inplace(x1, residual, weight, epsilon)
|
||||
|
||||
# FX graph after Dynamo tracing
|
||||
x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None
|
||||
out = torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace(x1, residual, weight, 1e-5); x1 = residual = None
|
||||
x2 = out[0]
|
||||
residual_out = out[1]
|
||||
```
|
||||
|
||||
### 2. AOTAutograd and Functionalization
|
||||
|
||||
AOTAutograd functionalizes the graph, converting any mutating operations to functional equivalents.
|
||||
For vLLM IR operations with `maybe_inplace` overloads, we perform this manually before AOTAutograd,
|
||||
converting them to the functional `default` overload using the pre-grad custom pass hook.
|
||||
|
||||
```python
|
||||
# After functionalization
|
||||
x1 = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5); x = None
|
||||
out = torch.ops.vllm_ir.fused_add_rms_norm.default(x1, residual, weight, 1e-5); x1 = residual = None
|
||||
x2 = out[0]
|
||||
residual_out = out[1]
|
||||
```
|
||||
|
||||
The pass also tracks which inputs were "donated" (passed to `maybe_inplace`),
|
||||
storing this information in vLLM's `PassContext` for later use in clone elimination.
|
||||
|
||||
### 3. IR Fusion and Transformation Passes
|
||||
|
||||
After functionalization, custom vLLM passes operate on the functional FX graph containing high-level IR operations.
|
||||
These passes can perform fusion, distribute operations for sequence parallelism, and other transformations:
|
||||
|
||||
```python
|
||||
# Example: Sequence Parallelism (see SequenceParallelismPass)
|
||||
# Before SP pass
|
||||
|
||||
all_reduce = torch.ops.vllm.all_reduce(x, "tp:0")
|
||||
rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5)
|
||||
|
||||
# after SP pass
|
||||
reduce_scatter = torch.ops.vllm.reduce_scatter(x, "tp:0")
|
||||
rms_norm = torch.ops.vllm_ir.rms_norm(all_reduce, weight, 1e-5)
|
||||
all_gather = torch.ops.vllm.all_gather(x, "tp:0")
|
||||
```
|
||||
|
||||
Fusion passes benefit from the high-level representation: they don't need to match against low-level PyTorch operations,
|
||||
handle different kernel implementations separately, or deal with functionalization of custom kernels.
|
||||
|
||||
### 4. IR Lowering
|
||||
|
||||
The lowering pass (`VllmIRLoweringPass`) replaces each vLLM IR operation with its selected implementation.
|
||||
The implementation is chosen based on the priority list and support predicates,
|
||||
using the **fake tensors** in the graph's metadata in place of op arguments:
|
||||
|
||||
```python
|
||||
# Implementation selection, same in eager dispatch and compile lowering
|
||||
def dispatch(*args) -> IrOpImpl:
|
||||
for provider in priority_list: # e.g., ["vllm_c", "native"]
|
||||
impl = ir_op.impls[provider]
|
||||
if not impl.supported:
|
||||
continue
|
||||
if impl.supports_args and not impl.supports_args(*args):
|
||||
continue
|
||||
return impl
|
||||
|
||||
# make_fx uses torch.fx.symbolic_trace
|
||||
impl_graph = make_fx(selected_impl.impl_fn)
|
||||
# Replace IR op node with impl_graph's nodes
|
||||
match.replace_by_example(selected_impl.impl_fn, node.args)
|
||||
```
|
||||
|
||||
For example, lowering `rms_norm` with the `vllm_c` implementation:
|
||||
|
||||
```python
|
||||
# Before lowering (IR op)
|
||||
rms_norm = torch.ops.vllm_ir.rms_norm.default(x, weight, 1e-5)
|
||||
|
||||
# After lowering (vllm_c implementation traced)
|
||||
# Note: Lowering does not currently functionalize, this will likely change in the future.
|
||||
empty = torch.ops.aten.empty.memory_format(x.shape, ...)
|
||||
rms_norm = torch.ops._C.rms_norm(empty, x, weight, 1e-5)
|
||||
```
|
||||
|
||||
When lowering an implementation that mutates inputs (`inplace=True`),
|
||||
the lowering pass inserts clones to preserve functional semantics:
|
||||
|
||||
```python
|
||||
# vllm_c implementation for fused_add_rms_norm mutates its first two arguments
|
||||
# Lowered with clones for safety
|
||||
clone_default = torch.ops.aten.clone.default(x)
|
||||
clone_default_1 = torch.ops.aten.clone.default(residual)
|
||||
fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(clone_default, clone_default_1, weight, 1e-5)
|
||||
```
|
||||
|
||||
### 5. Clone Cleanup
|
||||
|
||||
After lowering, the clone elimination pass (`UnsafeCloneEliminationPass`) removes unnecessary clones introduced during lowering.
|
||||
This pass is essential for achieving zero-copy behavior when using in-place kernels with `maybe_inplace`.
|
||||
The pass removes a clone if:
|
||||
|
||||
- the cloned input is created in the graph and not used again in the graph
|
||||
- the cloned input is a graph parameter, marked as donated
|
||||
|
||||
```python
|
||||
# After cleanup (donated inputs, no subsequent uses)
|
||||
fused_add_rms_norm = torch.ops._C.fused_add_rms_norm.default(x, residual, weight, 1e-5)
|
||||
```
|
||||
|
||||
The combination of inplace functionalization (tracking donated inputs) and clone cleanup enables the compiler to safely
|
||||
use in-place kernels without adding redundant copies or increasing the memory usage.
|
||||
|
||||
### 6. Inductor Optimization and Codegen
|
||||
|
||||
After IR lowering and cleanup, the graph contains only standard PyTorch operations and platform-specific custom ops.
|
||||
Inductor then performs its standard codegen:
|
||||
|
||||
- **Inductor lowering and pointwise fusion**: Fusing element-wise operations, reductions, etc.
|
||||
- **Memory planning**: Determining buffer allocation and reuse
|
||||
- **Kernel generation**: Generating Triton or C++ code for fused operations
|
||||
- **Autotuning**: Selecting the best kernel configurations
|
||||
|
||||
### Pipeline Summary
|
||||
|
||||
```text
|
||||
Model Forward Pass
|
||||
↓
|
||||
[Dynamo Tracing] → FX Graph with vllm_ir.* ops
|
||||
↓
|
||||
[Pre-grad: Inplace Functionalization] → maybe_inplace → default, track donated inputs
|
||||
↓
|
||||
[AOTAutograd] → Functionalization
|
||||
↓
|
||||
[Post-grad: IR Fusion Passes] → Fuse high-level IR ops (e.g., rms_norm + quant)
|
||||
↓
|
||||
[Post-grad: IR Lowering] → vllm_ir.* ops → impl ops (with clones if needed)
|
||||
↓
|
||||
[Post-grad: Clone Cleanup] → Remove unnecessary clones using donated input info
|
||||
↓
|
||||
[Inductor] → Pattern matching, fusion, memory planning, codegen
|
||||
↓
|
||||
Compiled Code
|
||||
```
|
||||
|
||||
## Core vLLM IR Concepts
|
||||
|
||||
### Operation Declaration
|
||||
|
||||
Operations are declared with the `@register_op` decorator, which creates an `IrOp` object:
|
||||
|
||||
```python
|
||||
@register_op(
|
||||
name=None, # Operation name (defaults to function name)
|
||||
activations=None, # List of activation parameters (defaults to params starting with 'x')
|
||||
allow_inplace=False, # Whether to create a maybe_inplace overload
|
||||
)
|
||||
def op_name(...):
|
||||
...
|
||||
```
|
||||
|
||||
**Parameters:**
|
||||
|
||||
- `activations`: List of parameter names considered "activations" (typically consumed by `maybe_inplace`). Defaults to parameters starting with `x`.
|
||||
- `allow_inplace`: Creates a `maybe_inplace` overload for memory-efficient execution (see below).
|
||||
|
||||
### The `maybe_inplace` Overload
|
||||
|
||||
The `maybe_inplace` overload is a critical feature for memory efficiency in LLM inference.
|
||||
It signals that the caller doesn't need to preserve the activation inputs after the operation,
|
||||
allowing in-place implementations to reuse input memory for outputs.
|
||||
|
||||
#### Semantics and Usage
|
||||
|
||||
```python
|
||||
# Standard usage: inputs are preserved
|
||||
out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon)
|
||||
# x and residual are unchanged, out and res_out are new tensors
|
||||
|
||||
# maybe_inplace: inputs may be modified
|
||||
out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon)
|
||||
# x and residual may be modified (undefined behavior to use them after this)
|
||||
# out and res_out may alias x and residual
|
||||
```
|
||||
|
||||
Using an activation input after passing it to `maybe_inplace` is **undefined behavior**:
|
||||
|
||||
```python
|
||||
# WRONG: Using x after donating it
|
||||
out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x, residual, weight, epsilon)
|
||||
result = out + x # ERROR: x was donated!
|
||||
```
|
||||
|
||||
If you need to preserve an input, either use the default overload or clone manually:
|
||||
|
||||
```python
|
||||
# Option 1: Use default overload
|
||||
out, res_out = ir.ops.fused_add_rms_norm(x, residual, weight, epsilon)
|
||||
result = out + x # OK: x is preserved
|
||||
|
||||
# Option 2: Clone before maybe_inplace
|
||||
out, res_out = ir.ops.fused_add_rms_norm.maybe_inplace(x.clone(), residual, weight, epsilon)
|
||||
result = out + x # OK: x is preserved, clone was donated
|
||||
```
|
||||
|
||||
#### Compilation Behavior
|
||||
|
||||
During compilation, the inplace functionalization pass validates that donated inputs are
|
||||
not used again and converts `maybe_inplace` to the functional `default` overload:
|
||||
|
||||
```python
|
||||
# Inplace functionalization pass (pre-grad)
|
||||
for node in graph.nodes:
|
||||
if node.target == torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace:
|
||||
# Check that activation inputs aren't used after this node
|
||||
for activation_arg in activation_inputs:
|
||||
for user in activation_arg.users:
|
||||
if user appears after node:
|
||||
raise ValueError(f"Input {activation_arg} donated but used again")
|
||||
|
||||
# Convert to default overload
|
||||
node.target = torch.ops.vllm_ir.fused_add_rms_norm.default
|
||||
|
||||
# Track donated graph inputs for later clone elimination
|
||||
for i, arg in enumerate(node.args):
|
||||
if arg.op == "placeholder" and i in activation_indices:
|
||||
pass_context.donated_input_ids.add(node_to_idx[arg])
|
||||
```
|
||||
|
||||
The donated input information is then used by the clone cleanup pass to eliminate
|
||||
unnecessary copies when in-place kernels are lowered.
|
||||
|
||||
#### Eager Mode Behavior
|
||||
|
||||
In eager mode (without `torch.compile`), `maybe_inplace` enables **maximally memory-efficient**
|
||||
execution by allowing the IR operation to dispatch directly to in-place implementations:
|
||||
|
||||
```python
|
||||
# Eager dispatch logic for maybe_inplace
|
||||
impl: IrOpImpl = ir_op.dispatch(*args)
|
||||
return impl.impl_fn(*args)
|
||||
|
||||
# Eager dispatch logic for default:
|
||||
impl: IrOpImpl = ir_op.dispatch(*args)
|
||||
if impl.inplace:
|
||||
args = [
|
||||
arg.clone() if i in ir_op.activations else arg
|
||||
for i, arg in enumerate(args)
|
||||
]
|
||||
return impl.impl_fn(*args)
|
||||
```
|
||||
|
||||
The combination of `maybe_inplace` in model code and in-place kernel implementations provides optimal memory efficiency
|
||||
in both eager and compiled modes, with identical semantics in both cases.
|
||||
|
||||
#### Memory Savings Example
|
||||
|
||||
Consider a transformer layer with residual connections:
|
||||
|
||||
```python
|
||||
# Without maybe_inplace (2 allocations per layer)
|
||||
hidden_states = self.attention(input)
|
||||
normed, residual = ir.ops.fused_add_rms_norm(hidden_states, input, weight, eps)
|
||||
# Memory: input (preserved), hidden_states (preserved), normed (new), residual (new)
|
||||
|
||||
# With maybe_inplace (0 allocations per layer when using in-place kernel)
|
||||
hidden_states = self.attention(input)
|
||||
normed, residual = ir.ops.fused_add_rms_norm.maybe_inplace(hidden_states, input, weight, eps)
|
||||
# Memory: normed (reuses hidden_states), residual (reuses input)
|
||||
```
|
||||
|
||||
### Implementation Registration
|
||||
|
||||
Implementations are registered using the `register_impl` method:
|
||||
|
||||
```python
|
||||
@ir.ops.op_name.register_impl(
|
||||
provider="provider_name", # Unique identifier (e.g., "vllm_c", "aiter", "triton")
|
||||
supported=True, # Static availability check
|
||||
supports_args=None, # Dynamic argument support check
|
||||
)
|
||||
def impl_fn(...):
|
||||
...
|
||||
```
|
||||
|
||||
**Provider naming conventions:**
|
||||
|
||||
- `native`: Reserved for the native torch implementation (declared with `@register_op`)
|
||||
- `vllm_c`: C++/CUDA kernels via `torch.ops._C`
|
||||
- `aiter`: AMD AITER library
|
||||
- `xpu_kernels`: SYCL/SYCLTLA kernels implemented in `vllm-xpu-kernels`
|
||||
- `triton_*`: Triton kernels
|
||||
- Platform/library names for other implementations
|
||||
|
||||
**Support checking:**
|
||||
|
||||
- `supported`: Static boolean, checked once at import time (e.g., `HAS_TRITON`, `is_cuda_alike()`)
|
||||
- `supports_args`: Function `(*args, **kwargs) -> bool` checking argument compatibility
|
||||
- Called with **fake tensors** during compilation for zero-cost checking
|
||||
- Called with **real tensors** during eager mode dispatch
|
||||
- Should NOT check batch sizes or add guards based on values
|
||||
|
||||
Example support predicate:
|
||||
|
||||
```python
|
||||
def aiter_rms_norm_supports(x, weight, epsilon, variance_size=None):
|
||||
# Check dtype (OK: doesn't depend on batch size)
|
||||
if x.dtype not in [torch.float16, torch.bfloat16]:
|
||||
return False
|
||||
# Check optional parameter (OK: static check)
|
||||
if variance_size is not None:
|
||||
return False
|
||||
return True
|
||||
|
||||
@ir.ops.rms_norm.register_impl("aiter", supports_args=aiter_rms_norm_supports)
|
||||
def rms_norm(...):
|
||||
...
|
||||
```
|
||||
|
||||
Batch-invariant kernels are automatically selected when `VLLM_BATCH_INVARIANT=1` is set.
|
||||
|
||||
### Eager Mode vs Compile Mode
|
||||
|
||||
vLLM IR operations behave identically in eager and compile modes:
|
||||
|
||||
**Eager mode:**
|
||||
|
||||
- Direct dispatch to implementation based on priority list
|
||||
- Support checked with real tensor arguments
|
||||
- Minimal overhead (can be optimized further if needed)
|
||||
|
||||
**Compile mode:**
|
||||
|
||||
- IR ops appear in FX graph as `torch.ops.vllm_ir.*` custom ops
|
||||
- Lowering selects implementation using fake tensors
|
||||
- Full integration with Inductor optimizations
|
||||
|
||||
This consistency enables:
|
||||
|
||||
- Prototyping in eager mode with confidence
|
||||
- Debugging by disabling compilation
|
||||
- Gradual migration from eager to compiled execution
|
||||
|
||||
## Other Topics
|
||||
|
||||
### Out-of-Tree Implementations
|
||||
|
||||
External platforms can register implementations without modifying vLLM:
|
||||
|
||||
```python
|
||||
# In external package
|
||||
from vllm import ir
|
||||
|
||||
@ir.ops.rms_norm.register_impl("my_platform", supported=is_my_platform())
|
||||
def rms_norm(x, weight, epsilon, variance_size=None):
|
||||
return my_platform.rms_norm(x, weight, epsilon)
|
||||
```
|
||||
|
||||
Then configure priority to use your implementation:
|
||||
|
||||
```python
|
||||
class MyPlatform(Platform):
|
||||
def get_default_ir_op_priority(self):
|
||||
return IrOpPriorityConfig(rms_norm=['my_platform', 'native'])
|
||||
|
||||
# Users can still override priority in the same way
|
||||
llm = LLM(ir_op_priority=IrOpPriorityConfig(rms_norm=['custom_oot_kernel']))
|
||||
```
|
||||
|
||||
### Debugging and Observability
|
||||
|
||||
!!! note
|
||||
Please let us know how observability can be improved for your use-case!
|
||||
|
||||
Enable debug logging to see kernel selection:
|
||||
|
||||
```bash
|
||||
VLLM_LOGGING_LEVEL=DEBUG vllm serve ...
|
||||
```
|
||||
|
||||
This logs:
|
||||
|
||||
- Which implementations are selected for each operation
|
||||
- Why implementations were rejected (unsupported, args not supported)
|
||||
- Compilation cache hits/misses
|
||||
- IR lowering statistics
|
||||
|
||||
Check selected implementations in compiled graphs:
|
||||
|
||||
```python
|
||||
# After compilation, inspect the lowering pass
|
||||
lowering_pass = backend.lowering_pass
|
||||
print(lowering_pass.selected_impls)
|
||||
# Output: {'rms_norm': {'node_123': 'vllm_c', 'node_456': 'vllm_c'}}
|
||||
```
|
||||
|
||||
## Migration from CustomOp
|
||||
|
||||
vLLM IR is designed to coexist with and gradually replace `CustomOp`:
|
||||
|
||||
1. **Op declaration**: Convert `CustomOp` class `PluggableLayer` and move `forward_native` to `@register_op` function
|
||||
2. **Implementation registration**: Use `@ir.ops.op_name.register_impl` instead of overriding methods
|
||||
3. **Layer usage**: Replace `self.op(...)` with `ir.ops.op_name(...)`
|
||||
4. **Configuration**: Migrate `--compilation-config.custom-ops` to `--ir-op-priority`
|
||||
|
||||
The migration can be done incrementally, one operation at a time.
|
||||
|
||||
## See Also
|
||||
|
||||
- [torch.compile Integration](torch_compile.md) - General compilation infrastructure
|
||||
- [Fusions](fusions.md) - Custom fusion and transformation passes in vLLM
|
||||
- [Custom Operations](custom_op.md) - Legacy custom op system
|
||||
@@ -52,10 +52,10 @@ th:not(:first-child) {
|
||||
| [mm](multimodal_inputs.md) | ✅ | ✅ | [🟠](https://github.com/vllm-project/vllm/pull/4194)<sup>^</sup> | ❔ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❔ | ✅ | | | |
|
||||
| best-of | ✅ | ✅ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](https://github.com/vllm-project/vllm/issues/7968) | ✅ | ✅ | | |
|
||||
| beam-search | ✅ | ✅ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/6137) | ✅ | ❌ | ✅ | ✅ | ✅ | ❔ | [❌](https://github.com/vllm-project/vllm/issues/7968) | ❔ | ✅ | ✅ | |
|
||||
| [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❔ | ❔ | ❌ | ❔ | ❔ | ✅ |
|
||||
| [prompt-embeds](prompt_embeds.md) | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❔ | ❔ | ✅ | ❔ | ❔ | ✅ |
|
||||
|
||||
\* Chunked prefill and prefix caching are only applicable to last-token or all pooling with causal attention.
|
||||
<sup>^</sup> LoRA is only applicable to the language backbone of multimodal models.
|
||||
<sup>^</sup> LoRA is only applicable to the language backbone of multimodal models.
|
||||
|
||||
### Feature x Hardware
|
||||
|
||||
|
||||
@@ -105,7 +105,7 @@ Batch invariance has been tested and verified on the following models:
|
||||
|
||||
- **DeepSeek series**: `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`
|
||||
- **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ`
|
||||
- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`
|
||||
- **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`, `Qwen/Qwen3-30B-A3B-Thinking-2507-FP8`
|
||||
- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
|
||||
- **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`
|
||||
- **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b`
|
||||
|
||||
@@ -36,10 +36,10 @@ The current reference pathway is **ExampleConnector**.
|
||||
Below ready-to-run scripts shows the workflow:
|
||||
|
||||
1 Encoder instance + 1 PD instance:
|
||||
`examples/online_serving/disaggregated_encoder/disagg_1e1pd_example.sh`
|
||||
`examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh`
|
||||
|
||||
1 Encoder instance + 1 Prefill instance + 1 Decode instance:
|
||||
`examples/online_serving/disaggregated_encoder/disagg_1e1p1d_example.sh`
|
||||
`examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh`
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -17,15 +17,15 @@ Two main reasons:
|
||||
|
||||
## Usage example
|
||||
|
||||
Please refer to [examples/online_serving/disaggregated_prefill.sh](../../examples/online_serving/disaggregated_prefill.sh) for the example usage of disaggregated prefilling.
|
||||
Please refer to [examples/disaggregated/disaggregated_prefill.sh](../../examples/disaggregated/disaggregated_prefill.sh) for the example usage of disaggregated prefilling.
|
||||
|
||||
Now supports 6 types of connectors:
|
||||
|
||||
- **ExampleConnector**: refer to [examples/offline_inference/disaggregated-prefill-v1/run.sh](../../examples/offline_inference/disaggregated-prefill-v1/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
|
||||
- **LMCacheConnectorV1**: refer to [examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/others/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
|
||||
- **ExampleConnector**: refer to [examples/disaggregated/example_connector/run.sh](../../examples/disaggregated/example_connector/run.sh) for the example usage of ExampleConnector disaggregated prefilling.
|
||||
- **LMCacheConnectorV1**: refer to [examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh](../../examples/disaggregated/lmcache/disagg_prefill_lmcache_v1/disagg_example_nixl.sh) for the example usage of LMCacheConnectorV1 disaggregated prefilling which uses NIXL as the underlying KV transmission.
|
||||
- **NixlConnector**: refer to [tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh](../../tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh) for the example usage of NixlConnector disaggregated prefilling which support fully async send/recv. For detailed usage guide, see [NixlConnector Usage Guide](nixl_connector_usage.md). For feature compatibility details, see [NixlConnector Compatibility Matrix](nixl_connector_compatibility.md).
|
||||
- **P2pNcclConnector**: refer to [examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/online_serving/disaggregated_serving_p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
|
||||
- **MooncakeConnector**: refer to [examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh](../../examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh) for the example usage of ExampleConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
|
||||
- **P2pNcclConnector**: refer to [examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh](../../examples/disaggregated/p2p_nccl_xpyd/disagg_example_p2p_nccl_xpyd.sh) for the example usage of P2pNcclConnector disaggregated prefilling.
|
||||
- **MooncakeConnector**: refer to [examples/disaggregated/mooncake_connector/run_mooncake_connector.sh](../../examples/disaggregated/mooncake_connector/run_mooncake_connector.sh) for the example usage of MooncakeConnector disaggregated prefilling. For detailed usage guide, see [MooncakeConnector Usage Guide](mooncake_connector_usage.md).
|
||||
- **MultiConnector**: take advantage of the kv_connector_extra_config: dict[str, Any] already present in KVTransferConfig to stash all the connectors we want in an ordered list of kwargs.such as:
|
||||
|
||||
```bash
|
||||
@@ -44,7 +44,7 @@ For NixlConnector, you may also specify one or multiple NIXL_Backend. Such as:
|
||||
--kv-transfer-config '{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"block_size": 64, "cpu_bytes_to_use": 1000000000}}'
|
||||
```
|
||||
|
||||
- **FlexKVConnectorV1**: refer to [examples/offline_inference/prefix_caching_flexkv.py](../../examples/offline_inference/prefix_caching_flexkv.py) for the example usage of FlexKVConnectorV1. FlexKV is a distributed KV Store and multi-level cache management system for ultra-large-scale LLM inference.
|
||||
- **FlexKVConnectorV1**: refer to [examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py](../../examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py) for the example usage of FlexKVConnectorV1. FlexKV is a distributed KV Store and multi-level cache management system for ultra-large-scale LLM inference.
|
||||
|
||||
```bash
|
||||
--kv-transfer-config '{"kv_connector":"FlexKVConnectorV1","kv_role":"kv_both"}'
|
||||
|
||||
@@ -31,7 +31,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct --port 8020 --kv-transfer-config '{"kv_conne
|
||||
### Proxy
|
||||
|
||||
```bash
|
||||
python examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020
|
||||
python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py --prefill http://192.168.0.2:8010 --decode http://192.168.0.3:8020
|
||||
```
|
||||
|
||||
Now you can send requests to the proxy server through port 8000.
|
||||
@@ -65,5 +65,5 @@ Now you can send requests to the proxy server through port 8000.
|
||||
|
||||
Refer to these example scripts in the vLLM repository:
|
||||
|
||||
- [run_mooncake_connector.sh](../../examples/online_serving/disaggregated_serving/mooncake_connector/run_mooncake_connector.sh)
|
||||
- [mooncake_connector_proxy.py](../../examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py)
|
||||
- [run_mooncake_connector.sh](../../examples/disaggregated/mooncake_connector/run_mooncake_connector.sh)
|
||||
- [mooncake_connector_proxy.py](../../examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py)
|
||||
|
||||
@@ -215,6 +215,67 @@ When loading RGBA images (images with transparency), vLLM converts them to RGB f
|
||||
- This setting only affects RGBA images with transparency; RGB images are unchanged
|
||||
- If not specified, the default white background `(255, 255, 255)` is used for backward compatibility
|
||||
|
||||
#### Moondream3 Prompt Recipes { #moondream3-prompt-recipes }
|
||||
|
||||
`Moondream3ForCausalLM` supports two task-specific prompt formats:
|
||||
|
||||
- `query`: ask a question about the image.
|
||||
- `caption`: generate a caption for the image.
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.assets.image import ImageAsset
|
||||
|
||||
llm = LLM(
|
||||
model="moondream/moondream3-preview",
|
||||
tokenizer="moondream/starmie-v1",
|
||||
trust_remote_code=True,
|
||||
max_model_len=2048,
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
|
||||
image = ImageAsset("stop_sign").pil_image
|
||||
|
||||
|
||||
def make_query_prompt(question: str) -> str:
|
||||
return (
|
||||
"<|endoftext|><image><|md_reserved_0|>query<|md_reserved_1|>"
|
||||
f"{question}<|md_reserved_2|>"
|
||||
)
|
||||
|
||||
|
||||
def make_caption_prompt(length: str = "normal") -> str:
|
||||
return (
|
||||
"<|endoftext|><image><|md_reserved_0|>"
|
||||
f"describe<|md_reserved_1|>{length}<|md_reserved_2|>"
|
||||
)
|
||||
|
||||
|
||||
query_out = llm.generate(
|
||||
{
|
||||
"prompt": make_query_prompt("What is shown in this image?"),
|
||||
"multi_modal_data": {"image": image},
|
||||
},
|
||||
SamplingParams(max_tokens=64, temperature=0),
|
||||
)[0].outputs[0].text
|
||||
|
||||
caption_out = llm.generate(
|
||||
{
|
||||
"prompt": make_caption_prompt(),
|
||||
"multi_modal_data": {"image": image},
|
||||
},
|
||||
SamplingParams(max_tokens=100, temperature=0),
|
||||
)[0].outputs[0].text
|
||||
|
||||
print("query:", query_out)
|
||||
print("caption:", caption_out)
|
||||
```
|
||||
|
||||
!!! note
|
||||
The native Moondream3 model also has `detect` and `point` skills. Those
|
||||
require custom coordinate decoding and are not exposed by this vLLM
|
||||
implementation.
|
||||
|
||||
### Video Inputs
|
||||
|
||||
You can pass a list of NumPy arrays directly to the `'video'` field of the multi-modal dictionary
|
||||
|
||||
@@ -20,12 +20,47 @@ You can pass prompt embeddings from Hugging Face Transformers models to the `'p
|
||||
|
||||
## Online Serving
|
||||
|
||||
Our OpenAI-compatible server accepts prompt embeddings inputs via the [Completions API](https://platform.openai.com/docs/api-reference/completions). Prompt embeddings inputs are added via a new `'prompt_embeds'` key in the JSON package and are enabled by the `--enable-prompt-embeds` flag in `vllm serve`.
|
||||
Our OpenAI-compatible server accepts prompt embeddings inputs via both the [Completions API](https://platform.openai.com/docs/api-reference/completions) and the [Chat Completions API](https://platform.openai.com/docs/api-reference/chat). Both are enabled by the `--enable-prompt-embeds` flag in `vllm serve`.
|
||||
|
||||
### Completions API
|
||||
|
||||
Prompt embeddings inputs are added via a `'prompt_embeds'` key in the JSON request body.
|
||||
|
||||
When a mixture of `'prompt_embeds'` and `'prompt'` inputs are provided in a single request, the prompt embeds are always returned first.
|
||||
|
||||
Prompt embeddings are passed in as base64 encoded torch tensors.
|
||||
|
||||
The Completions endpoint does **not** apply a chat template to `prompt_embeds`. If the model assumes some chat template, the caller is responsible for producing embeddings for the full, already-templated prompt: apply the chat template, then embed the resulting token IDs. Anything the model would normally need (system prompt, role markers, generation prompt, etc.) must already be baked into the embedded tokens.
|
||||
|
||||
### Chat Completions API
|
||||
|
||||
Prompt embeddings can be included as content parts in chat messages, interleaved with text:
|
||||
|
||||
```json
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": "You are a helpful assistant."},
|
||||
{"type": "prompt_embeds", "data": "<base64_encoded_tensor>"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "prompt_embeds", "data": "<base64_encoded_tensor>"},
|
||||
{"type": "text", "text": "Summarize the above."}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
Each `prompt_embeds` content part contains a `data` field with a base64-encoded `torch.Tensor` of shape `(num_tokens, hidden_size)`. Multiple `prompt_embeds` parts can appear in any message, in any position relative to text parts. The server expands each part into the correct number of placeholder tokens during chat template rendering, then splices the pre-computed embeddings into the model's input at the corresponding positions.
|
||||
|
||||
Unlike the Completions API, a `prompt_embeds` content part should encode **only** the content, not a templated conversation. The server wraps the chat template around the embedded content at request time, the same way it would for a plain text `content` string. Embedding a full templated conversation here would double-apply the template and produce incorrect inputs to the model.
|
||||
|
||||
!!! warning
|
||||
The vLLM engine may crash if incorrect shape of embeddings is passed.
|
||||
Only enable this flag for trusted users!
|
||||
|
||||
@@ -7,7 +7,7 @@ import sys
|
||||
import textwrap
|
||||
import traceback
|
||||
from argparse import SUPPRESS, Action, HelpFormatter
|
||||
from collections.abc import Iterable
|
||||
from collections.abc import Callable, Iterable
|
||||
from importlib.machinery import ModuleSpec
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Literal
|
||||
@@ -48,6 +48,7 @@ class MockPluggableLayer:
|
||||
|
||||
|
||||
mock_if_no_torch("vllm._C", MagicMock())
|
||||
mock_if_no_torch("vllm._C_stable_libtorch", MagicMock())
|
||||
mock_if_no_torch(
|
||||
"vllm.model_executor.custom_op",
|
||||
MagicMock(CustomOp=MockCustomOp, PluggableLayer=MockPluggableLayer),
|
||||
@@ -67,6 +68,31 @@ importlib.metadata.version = lambda name: VERSIONS.get(name) or "0.0.0"
|
||||
mock_if_no_torch("torch.nn", MagicMock(Parameter=object))
|
||||
|
||||
|
||||
# Mock torch.library.infer_schema for vllm.ir.ops.IrOpInplaceOverload.__init__
|
||||
# We need to return the corresponding number of inputs, as IR infra will assert it
|
||||
def get_outputs(native_fn: Callable) -> str:
|
||||
"""
|
||||
Extract output schema from function's return type annotation,
|
||||
e.g. 'Tensor' or 'Tensor, Tensor'.
|
||||
"""
|
||||
import typing
|
||||
|
||||
return_type = typing.get_type_hints(native_fn)["return"]
|
||||
origin = typing.get_origin(return_type)
|
||||
arg_name = lambda a: a.__name__ if hasattr(a, "__name__") else str(a)
|
||||
if origin is tuple:
|
||||
args = typing.get_args(return_type)
|
||||
return ", ".join(arg_name(arg) for arg in args)
|
||||
else:
|
||||
return f"{arg_name(return_type)}"
|
||||
|
||||
|
||||
mock_if_no_torch(
|
||||
"torch.library",
|
||||
MagicMock(infer_schema=lambda fn, **k: f"(Tensor x) -> {get_outputs(fn)}"),
|
||||
)
|
||||
|
||||
|
||||
class PydanticMagicMock(MagicMock):
|
||||
"""`MagicMock` that's able to generate pydantic-core schemas."""
|
||||
|
||||
|
||||
@@ -599,6 +599,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I<sup>+</sup> | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ |
|
||||
| `MolmoForCausalLM` | Molmo | T + I<sup>+</sup> | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ |
|
||||
| `Molmo2ForConditionalGeneration` | Molmo2 | T + I<sup>+</sup> / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B` | ✅︎ | ✅︎ |
|
||||
| `Moondream3ForCausalLM` | Moondream3 | T + I | `moondream/moondream3-preview` | | ✅︎ |
|
||||
| `MusicFlamingoForConditionalGeneration` | MusicFlamingo | T + A | `nvidia/music-flamingo-2601-hf`, `nvidia/music-flamingo-think-2601-hf` | ✅︎ | ✅︎ |
|
||||
| `NVLM_D_Model` | NVLM-D 1.0 | T + I<sup>+</sup> | `nvidia/NVLM-D-72B`, etc. | | ✅︎ |
|
||||
| `OpenCUAForConditionalGeneration` | OpenCUA-7B | T + I<sup>E+</sup> | `xlangai/OpenCUA-7B` | ✅︎ | ✅︎ |
|
||||
@@ -613,6 +614,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `Phi4MMForCausalLM` | Phi-4-multimodal | T + I<sup>+</sup> / T + A<sup>+</sup> / I<sup>+</sup> + A<sup>+</sup> | `microsoft/Phi-4-multimodal-instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `Phi4ForCausalLMV` | Phi-4-reasoning-vision | T + I<sup>+</sup> | `microsoft/Phi-4-reasoning-vision-15B`, etc. | | ✅︎ |
|
||||
| `PixtralForConditionalGeneration` | Ministral 3 (Mistral format), Mistral 3 (Mistral format), Mistral Large 3 (Mistral format), Pixtral (Mistral format) | T + I<sup>+</sup> | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, `mistralai/Mistral-Large-3-675B-Instruct-2512` `mistralai/Pixtral-12B-2409` etc. | ✅︎ | ✅︎ |
|
||||
| `QianfanOCRForConditionalGeneration` | QianfanOCR | T + I<sup>E+</sup> | `baidu/Qianfan-OCR`, etc. | ✅︎ | ✅︎ |
|
||||
| `QwenVLForConditionalGeneration`<sup>^</sup> | Qwen-VL | T + I<sup>E+</sup> | `Qwen/Qwen-VL`, `Qwen/Qwen-VL-Chat`, etc. | ✅︎ | ✅︎ |
|
||||
| `Qwen2AudioForConditionalGeneration` | Qwen2-Audio | T + A<sup>+</sup> | `Qwen/Qwen2-Audio-7B-Instruct` | | ✅︎ |
|
||||
| `Qwen2VLForConditionalGeneration` | QVQ, Qwen2-VL | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/QVQ-72B-Preview`, `Qwen/Qwen2-VL-7B-Instruct`, `Qwen/Qwen2-VL-72B-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -661,6 +663,12 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
!!! note
|
||||
For `InternVLChatModel`, only InternVL2.5 with Qwen2.5 text backbone (`OpenGVLab/InternVL2.5-1B` etc.), InternVL3 and InternVL3.5 have video inputs support currently.
|
||||
|
||||
!!! note
|
||||
`Moondream3ForCausalLM` uses task-specific prompt templates for `query`
|
||||
and `caption`. The native `detect` and `point` skills require custom
|
||||
coordinate decoding and are not exposed by this vLLM implementation.
|
||||
See [Moondream3 prompt recipes](../features/multimodal_inputs.md#moondream3-prompt-recipes).
|
||||
|
||||
!!! note
|
||||
To use `TIGER-Lab/Mantis-8B-siglip-llama3`, you have to pass `--hf_overrides '{"architectures": ["MantisForConditionalGeneration"]}'` when running vLLM.
|
||||
|
||||
|
||||
@@ -98,7 +98,7 @@ For larger scale deployments especially, it can make sense to handle the orchest
|
||||
|
||||
In this case, it's more convenient to treat each DP rank like a separate vLLM deployment, with its own endpoint, and have an external router balance HTTP requests between them, making use of appropriate real-time telemetry from each server for routing decisions.
|
||||
|
||||
This can already be done trivially for non-MoE models, since each deployed server is fully independent. No data parallel CLI options need to be used for this.
|
||||
This can already be done trivially for non-MoE models, since each deployed server is fully independent. In that case, launch independent vLLM instances without any `--data-parallel-*` arguments; external DP CLI options are only supported for MoE deployments.
|
||||
|
||||
We support an equivalent topology for MoE DP+EP which can be configured via the following CLI arguments.
|
||||
|
||||
|
||||
@@ -4,11 +4,11 @@ For general troubleshooting, see [Troubleshooting](../usage/troubleshooting.md).
|
||||
|
||||
## Verify inter-node GPU communication
|
||||
|
||||
After you start the Ray cluster, verify GPU-to-GPU communication across nodes. Proper configuration can be non-trivial. For more information, see [troubleshooting script](../usage/troubleshooting.md#incorrect-hardwaredriver). If you need additional environment variables for communication configuration, append them to [examples/online_serving/run_cluster.sh](../../examples/online_serving/run_cluster.sh), for example `-e NCCL_SOCKET_IFNAME=eth0`. Setting environment variables during cluster creation is recommended because the variables propagate to all nodes. In contrast, setting environment variables in the shell affects only the local node. For more information, see <https://github.com/vllm-project/vllm/issues/6803>.
|
||||
After you start the Ray cluster, verify GPU-to-GPU communication across nodes. Proper configuration can be non-trivial. For more information, see [troubleshooting script](../usage/troubleshooting.md#incorrect-hardwaredriver). If you need additional environment variables for communication configuration, append them to [examples/ray_serving/run_cluster.sh](../../examples/ray_serving/run_cluster.sh), for example `-e NCCL_SOCKET_IFNAME=eth0`. Setting environment variables during cluster creation is recommended because the variables propagate to all nodes. In contrast, setting environment variables in the shell affects only the local node. For more information, see <https://github.com/vllm-project/vllm/issues/6803>.
|
||||
|
||||
## No available node types can fulfill resource request
|
||||
|
||||
The error message `Error: No available node types can fulfill resource request` can appear even when the cluster has enough GPUs. The issue often occurs when nodes have multiple IP addresses and vLLM can't select the correct one. Ensure that vLLM and Ray use the same IP address by setting `VLLM_HOST_IP` in [examples/online_serving/run_cluster.sh](../../examples/online_serving/run_cluster.sh) (with a different value on each node). Use `ray status` and `ray list nodes` to verify the chosen IP address. For more information, see <https://github.com/vllm-project/vllm/issues/7815>.
|
||||
The error message `Error: No available node types can fulfill resource request` can appear even when the cluster has enough GPUs. The issue often occurs when nodes have multiple IP addresses and vLLM can't select the correct one. Ensure that vLLM and Ray use the same IP address by setting `VLLM_HOST_IP` in [examples/ray_serving/run_cluster.sh](../../examples/ray_serving/run_cluster.sh) (with a different value on each node). Use `ray status` and `ray list nodes` to verify the chosen IP address. For more information, see <https://github.com/vllm-project/vllm/issues/7815>.
|
||||
|
||||
## Ray observability
|
||||
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
# Codex
|
||||
|
||||
[Codex](https://github.com/openai/codex) is OpenAI's official agentic coding tool that lives in your terminal. It can understand your codebase, edit files, run commands, and help you write code more efficiently.
|
||||
|
||||
By pointing Codex at a vLLM server, you can use your own models as the backend instead of the OpenAI API. This is useful for:
|
||||
|
||||
- Running fully local/private coding assistance
|
||||
- Using open-weight models with tool calling capabilities
|
||||
- Testing and developing with custom models
|
||||
|
||||
## How It Works
|
||||
|
||||
vLLM implements the OpenAI-Responses API, which is the same API that Codex uses to communicate with OpenAI's servers. By configuring Codex to point at your vLLM server, Codex sends its requests to vLLM instead of OpenAI. vLLM then translates these requests to work with your local model and returns responses in the format Codex expects.
|
||||
|
||||
This means any model served by vLLM with proper tool calling support can act as a drop-in replacement for OpenAI models in Codex.
|
||||
|
||||
## Requirements
|
||||
|
||||
Codex requires a model with strong tool calling capabilities. The model must support the OpenAI-Responses tool calling API. See [Tool Calling](../../features/tool_calling.md) for details on enabling tool calling for your model.
|
||||
|
||||
## Installation
|
||||
|
||||
First, install Codex by following the [official installation guide](https://github.com/openai/codex).
|
||||
|
||||
## Starting the vLLM Server
|
||||
|
||||
Start vLLM with a tool-calling capable model - here's an example using `Qwen/Qwen3-27B`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3.6-27B --port 8000 --tensor-parallel-size 8 --max-model-len 262144 --reasoning-parser qwen3 --enable-auto-tool-choice --tool-call-parser qwen3_coder
|
||||
|
||||
```
|
||||
|
||||
For other models, you'll need to enable tool calling explicitly with `--enable-auto-tool-choice` and the right `--tool-call-parser`. Refer to the [Tool Calling documentation](../../features/tool_calling.md) for the correct flags for your model.
|
||||
|
||||
## Configuring Codex
|
||||
|
||||
Codex is configured via a TOML file located at `~/.codex/config.toml`. Create or edit this file to point Codex at your vLLM server:
|
||||
|
||||
```toml
|
||||
model = "my-model"
|
||||
model_provider = "vllm"
|
||||
|
||||
[model_providers.vllm]
|
||||
name = "vLLM"
|
||||
env_key = "VLLM_API_KEY"
|
||||
base_url = "http://localhost:8000/v1"
|
||||
wire_api = "responses"
|
||||
```
|
||||
|
||||
The configuration fields:
|
||||
|
||||
| Field | Description |
|
||||
| ----- | ----------- |
|
||||
| `model` | The model name to use. Must match the `--served-model-name` you passed to vLLM. |
|
||||
| `model_provider` | Set to `"vllm"` to use your local vLLM server. |
|
||||
| `[model_providers.vllm]` | Configuration section for the vLLM provider. |
|
||||
| `name` | A display name for your vLLM provider. |
|
||||
| `env_key` | The name of an environment variable that Codex will read for the API key. vLLM does not require authentication by default, so this can be any value. |
|
||||
| `base_url` | The URL of your vLLM server's OpenAI-compatible API endpoint (default is `http://localhost:8000/v1`). |
|
||||
| `wire_api` | The API style to use. Set to `"responses"` for the OpenAI Responses API |
|
||||
|
||||
!!! tip
|
||||
You can set the `env_key` to any dummy environment variable since vLLM doesn't require authentication by default:
|
||||
```bash
|
||||
export VLLM_API_KEY=dummy
|
||||
```
|
||||
|
||||
!!! warning
|
||||
When using the `responses` API, ensure your vLLM version supports the OpenAI Responses API.
|
||||
|
||||
## Testing the Setup
|
||||
|
||||
Once Codex is configured, launch it in your project directory:
|
||||
|
||||
```bash
|
||||
codex
|
||||
```
|
||||
|
||||
Try a simple prompt to verify the connection, such as asking it to explain a file in your project. If the model responds correctly, your setup is working. You can now use Codex with your vLLM-served model for coding tasks.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
**Connection refused**: Ensure vLLM is running and accessible at the specified URL. Check that the port matches and that `base_url` includes the `/v1` path suffix.
|
||||
|
||||
**Tool calls not working**: Verify that your model supports tool calling and that you've enabled it with the correct `--tool-call-parser` flag. See [Tool Calling](../../features/tool_calling.md).
|
||||
|
||||
**Model not found**: Ensure the `model` field in `~/.codex/config.toml` matches the `--served-model-name` you passed to vLLM.
|
||||
@@ -542,6 +542,6 @@ Key capabilities:
|
||||
- Scales from a single GPU to a multi-node cluster without code changes.
|
||||
- Provides observability and autoscaling policies through Ray dashboards and metrics.
|
||||
|
||||
The following example shows how to deploy a large model like DeepSeek R1 with Ray Serve LLM: [examples/online_serving/ray_serve_deepseek.py](../../examples/online_serving/ray_serve_deepseek.py).
|
||||
The following example shows how to deploy a large model like DeepSeek R1 with Ray Serve LLM: [examples/ray_serving/ray_serve_deepseek.py](../../examples/ray_serving/ray_serve_deepseek.py).
|
||||
|
||||
Learn more about Ray Serve LLM with the official [Ray Serve LLM documentation](https://docs.ray.io/en/latest/serve/llm/index.html).
|
||||
|
||||
@@ -78,7 +78,7 @@ For details, see the [Ray documentation](https://docs.ray.io/en/latest/index.htm
|
||||
|
||||
### Ray cluster setup with containers
|
||||
|
||||
The helper script [examples/online_serving/run_cluster.sh](../../examples/online_serving/run_cluster.sh) starts containers across nodes and initializes Ray. By default, the script runs Docker without administrative privileges, which prevents access to the GPU performance counters when profiling or tracing. To enable admin privileges, add the `--cap-add=CAP_SYS_ADMIN` flag to the Docker command.
|
||||
The helper script [examples/ray_serving/run_cluster.sh](../../examples/ray_serving/run_cluster.sh) starts containers across nodes and initializes Ray. By default, the script runs Docker without administrative privileges, which prevents access to the GPU performance counters when profiling or tracing. To enable admin privileges, add the `--cap-add=CAP_SYS_ADMIN` flag to the Docker command.
|
||||
|
||||
Choose one node as the head node and run:
|
||||
|
||||
@@ -162,7 +162,7 @@ vllm serve /path/to/the/model/in/the/container \
|
||||
|
||||
Efficient tensor parallelism requires fast internode communication, preferably through high-speed network adapters such as InfiniBand.
|
||||
To set up the cluster to use InfiniBand, append additional arguments like `--privileged -e NCCL_IB_HCA=mlx5` to the
|
||||
[examples/online_serving/run_cluster.sh](../../examples/online_serving/run_cluster.sh) helper script.
|
||||
[examples/ray_serving/run_cluster.sh](../../examples/ray_serving/run_cluster.sh) helper script.
|
||||
Contact your system administrator for more information about the required flags.
|
||||
|
||||
## Enabling GPUDirect RDMA
|
||||
|
||||
+1
-1
@@ -5,7 +5,7 @@ This file provides a disaggregated prefilling proxy demo to demonstrate an
|
||||
example usage of XpYd disaggregated prefilling.
|
||||
We can launch multiple vllm instances (2 for prefill and 2 for decode), and
|
||||
launch this proxy demo through:
|
||||
python3 examples/online_serving/disaggregated_serving/disagg_proxy_demo.py \
|
||||
python3 examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py \
|
||||
--model $model_name \
|
||||
--prefill localhost:8100 localhost:8101 \
|
||||
--decode localhost:8200 localhost:8201 \
|
||||
Executable → Regular
+1
-1
@@ -5,6 +5,6 @@ This example contains scripts that demonstrate disaggregated prefill in the offl
|
||||
## Files
|
||||
|
||||
- `run.sh` - A helper script that will run `prefill_example.py` and `decode_example.py` sequentially.
|
||||
- Make sure you are in the `examples/offline_inference/disaggregated-prefill-v1` directory before running `run.sh`.
|
||||
- Make sure you are in the `examples/disaggregated/example_connector` directory before running `run.sh`.
|
||||
- `prefill_example.py` - A script which performs prefill only, saving the KV state to the `local_storage` directory and the prompts to `output.txt`.
|
||||
- `decode_example.py` - A script which performs decode only, loading the KV state from the `local_storage` directory and the prompts from `output.txt`.
|
||||
+1
-1
@@ -14,7 +14,7 @@ Requirements:
|
||||
|
||||
Usage:
|
||||
1. Run this script:
|
||||
python examples/offline_inference/prefix_caching_flexkv.py \
|
||||
python examples/disaggregated/flexkv_connector/prefix_caching_flexkv.py \
|
||||
--model /path/to/your/model
|
||||
|
||||
2. Arguments:
|
||||
+2
-2
@@ -1,12 +1,12 @@
|
||||
# KV Load Failure Recovery Test
|
||||
|
||||
This example builds upon the `disaggregated-prefill-v1` example in `examples/offline_inference`.
|
||||
This example builds upon the `example_connector` example in `examples/disaggregated`.
|
||||
|
||||
It demonstrates vLLM's ability to recover from KV load failures in both synchronous and asynchronous loading modes. The goal is to verify that vLLM correctly identifies invalid KV blocks, reschedules the affected requests, and ensures successful and consistent output.
|
||||
|
||||
## Files
|
||||
|
||||
- `prefill_example.py` – performs the prefill stage and saves KV data (same as in `disaggregated-prefill-v1`).
|
||||
- `prefill_example.py` – performs the prefill stage and saves KV data (same as in `example_connector`).
|
||||
- `decode_example.py` – performs the decode stage. Accepts:
|
||||
- `--simulate-failure`: simulates KV load failure using a custom connector.
|
||||
- `--async-load`: enables asynchronous KV loading mode.
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user