forked from Karylab-cklius/vllm
Compare commits
105
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
77230471c0 | ||
|
|
efb4cdf2b8 | ||
|
|
92a7c121b6 | ||
|
|
307b17ce33 | ||
|
|
3ca6ca210f | ||
|
|
10558f5f46 | ||
|
|
121dbe7a22 | ||
|
|
f03d82efdd | ||
|
|
a7fb008510 | ||
|
|
ff449b6426 | ||
|
|
3527229517 | ||
|
|
b55b26520c | ||
|
|
3179e53135 | ||
|
|
efdc95674d | ||
|
|
54146a9bf9 | ||
|
|
ca97f7b9bb | ||
|
|
a04e0cf3b8 | ||
|
|
cb1b02d0e8 | ||
|
|
a749a33d8d | ||
|
|
c42981d034 | ||
|
|
0ff1bf9bb1 | ||
|
|
0ab67c0222 | ||
|
|
3795d7acf4 | ||
|
|
18599bfdf2 | ||
|
|
296741d025 | ||
|
|
a966aaed30 | ||
|
|
6841f5dc77 | ||
|
|
c2fb013312 | ||
|
|
ccfb620c62 | ||
|
|
0335316a9b | ||
|
|
944e138bcf | ||
|
|
b58669cb42 | ||
|
|
1628239eb2 | ||
|
|
93da1fe97a | ||
|
|
169988a3c0 | ||
|
|
faab189554 | ||
|
|
6f20f81cbf | ||
|
|
d1a75e303d | ||
|
|
4a42aba380 | ||
|
|
a80d6f150c | ||
|
|
91a2d39014 | ||
|
|
a05848e255 | ||
|
|
51fda1ba44 | ||
|
|
39a7f4f4e2 | ||
|
|
b92ef9ec5a | ||
|
|
5560cac7e2 | ||
|
|
5b39b268f5 | ||
|
|
22524f7a92 | ||
|
|
9d8ad5b408 | ||
|
|
11b69129e2 | ||
|
|
33f36d4260 | ||
|
|
37e288214b | ||
|
|
5371d6fb40 | ||
|
|
6d7d4da99e | ||
|
|
3f1a4bb639 | ||
|
|
762022cafb | ||
|
|
3885d340a4 | ||
|
|
ef70057ca7 | ||
|
|
e48cb85185 | ||
|
|
92879e12ba | ||
|
|
68dd7db810 | ||
|
|
8a8c9b564e | ||
|
|
a269744e9f | ||
|
|
8b49cf3a37 | ||
|
|
2ae73c758c | ||
|
|
d95d03c719 | ||
|
|
803b9d7881 | ||
|
|
1312f07531 | ||
|
|
fa1b9840f6 | ||
|
|
916e56c05c | ||
|
|
a085b5257d | ||
|
|
7fd05e05ae | ||
|
|
99255f3cb5 | ||
|
|
75a7cf2c10 | ||
|
|
4b95e9cec4 | ||
|
|
856b15c62c | ||
|
|
6fb3f7b46b | ||
|
|
d109eacd05 | ||
|
|
e68fa1b90a | ||
|
|
f05f3664c3 | ||
|
|
e9f8f31e9a | ||
|
|
de3fe8dc62 | ||
|
|
0899f436aa | ||
|
|
358a755e43 | ||
|
|
a60883644b | ||
|
|
5aa371dc8e | ||
|
|
de3da0b97c | ||
|
|
9e92de51c6 | ||
|
|
bde0efdbb7 | ||
|
|
ea74f701db | ||
|
|
a8208e6a81 | ||
|
|
76c9cccc36 | ||
|
|
ed57f77192 | ||
|
|
7a1eb8ac2e | ||
|
|
c2e88a281c | ||
|
|
fd74c90d9c | ||
|
|
146f44b77d | ||
|
|
0d4f714208 | ||
|
|
03aeed802f | ||
|
|
2c8b76c5cb | ||
|
|
407b34be26 | ||
|
|
4c7c69b4e0 | ||
|
|
5e2c37facd | ||
|
|
c8bbe05189 | ||
|
|
6232fb4b66 |
@@ -69,11 +69,11 @@ steps:
|
||||
pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
|
||||
pytest -x -v -s tests/quantization/test_cpu_wna16.py"
|
||||
|
||||
- label: CPU-Distributed Tests
|
||||
- label: CPU-Distributed Tests (PP+TP)
|
||||
depends_on: []
|
||||
device: intel_cpu
|
||||
no_plugin: true
|
||||
source_file_dependencies:
|
||||
source_file_dependencies: &cpu_distributed_deps
|
||||
- csrc/cpu/shm.cpp
|
||||
- vllm/v1/worker/cpu_worker.py
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
@@ -82,10 +82,21 @@ steps:
|
||||
- vllm/platforms/cpu.py
|
||||
- vllm/distributed/parallel_state.py
|
||||
- vllm/distributed/device_communicators/cpu_communicator.py
|
||||
- .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 10m "
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh"
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh tp_pp"
|
||||
|
||||
- label: CPU-Distributed Tests (DP+TP)
|
||||
depends_on: []
|
||||
device: intel_cpu
|
||||
no_plugin: true
|
||||
source_file_dependencies: *cpu_distributed_deps
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 10m "
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh dp_tp"
|
||||
|
||||
- label: CPU-Multi-Modal Model Tests %N
|
||||
depends_on: []
|
||||
|
||||
@@ -192,6 +192,7 @@ export BUILDKITE_COMMIT
|
||||
export PARENT_COMMIT
|
||||
export IMAGE_TAG
|
||||
export IMAGE_TAG_LATEST
|
||||
export COMMIT="${COMMIT:-${BUILDKITE_COMMIT}}"
|
||||
export CACHE_FROM
|
||||
export CACHE_FROM_BASE_BRANCH
|
||||
export CACHE_FROM_MAIN
|
||||
|
||||
@@ -126,5 +126,4 @@ steps:
|
||||
'cd tests &&
|
||||
pytest -v -s lora/test_default_mm_loras.py &&
|
||||
(pytest -v -s lora/test_qwen3_unembed.py || true) &&
|
||||
(pytest -v -s lora/test_qwenvl.py || true) &&
|
||||
pytest -v -s lora/test_whisper.py'
|
||||
|
||||
@@ -27,7 +27,7 @@ steps:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --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 manylinux_2_31"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -40,7 +40,7 @@ steps:
|
||||
- "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 ."
|
||||
- "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 manylinux_2_35"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -53,7 +53,7 @@ steps:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
|
||||
- "mkdir artifacts"
|
||||
- "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 manylinux_2_35"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -66,7 +66,7 @@ steps:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --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 manylinux_2_31"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -79,7 +79,7 @@ steps:
|
||||
- "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 ."
|
||||
- "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 manylinux_2_35"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -92,7 +92,7 @@ steps:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
|
||||
- "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 manylinux_2_35"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -121,7 +121,19 @@ steps:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
$(bash .buildkite/scripts/docker-build-metadata-args.sh) \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
# re-tag to default image tag and push, just in case arm64 build fails
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
@@ -134,7 +146,19 @@ steps:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
$(bash .buildkite/scripts/docker-build-metadata-args.sh) \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 12.9"
|
||||
@@ -144,7 +168,18 @@ steps:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=12.9.1 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
|
||||
# re-tag to default image tag and push, just in case arm64 build fails
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
@@ -157,7 +192,18 @@ steps:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=12.9.1 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
@@ -167,7 +213,21 @@ steps:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
$(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg UBUNTU_VERSION=24.04 \
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
@@ -179,7 +239,21 @@ steps:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
$(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg UBUNTU_VERSION=24.04 \
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
@@ -189,7 +263,20 @@ steps:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=12.9.1 \
|
||||
--build-arg UBUNTU_VERSION=24.04 \
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
@@ -201,7 +288,20 @@ steps:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64_CU129}\" --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- |
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg CUDA_VERSION=12.9.1 \
|
||||
--build-arg UBUNTU_VERSION=24.04 \
|
||||
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
|
||||
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
|
||||
--build-arg INSTALL_KV_CONNECTORS=true \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile .
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
|
||||
|
||||
- block: "Build release image for x86_64 CPU"
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Detect the manylinux platform tag for a wheel and rename it in place.
|
||||
|
||||
vLLM's build images produce wheels with the generic ``linux_<arch>`` platform
|
||||
tag, which installers like ``pip`` won't accept off PyPI/our index. We need to
|
||||
rewrite the platform tag to the appropriate ``manylinux_<major>_<minor>_<arch>``
|
||||
before uploading.
|
||||
|
||||
Historically the tag was hard-coded per build (``manylinux_2_31`` for the
|
||||
Ubuntu 20.04-based image, ``manylinux_2_35`` for the Ubuntu 22.04-based
|
||||
images). That is brittle: bumping the base image silently produces wheels
|
||||
labelled with the wrong glibc requirement. This script asks ``auditwheel``
|
||||
to derive the tag from the symbol versions actually referenced by the
|
||||
binaries inside the wheel, so the label tracks reality.
|
||||
|
||||
We can't simply call ``auditwheel repair`` -- it tries to graft external
|
||||
shared libraries into the wheel and fails on vLLM's CUDA/cuBLAS dependencies.
|
||||
Instead we use ``auditwheel.wheel_abi.analyze_wheel_abi`` directly, which is
|
||||
the same call that powers ``auditwheel show``, and read off
|
||||
``winfo.sym_policy.name``.
|
||||
|
||||
Usage:
|
||||
detect-manylinux-tag.py <wheel_path>
|
||||
|
||||
The wheel is renamed in place; the new path is printed on stdout. All
|
||||
diagnostics go to stderr so callers can capture stdout safely.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from auditwheel.error import (
|
||||
AuditwheelError,
|
||||
NonPlatformWheelError,
|
||||
WheelToolsError,
|
||||
)
|
||||
from auditwheel.wheel_abi import analyze_wheel_abi
|
||||
from auditwheel.wheeltools import get_wheel_architecture, get_wheel_libc
|
||||
|
||||
|
||||
def detect_platform_tag(wheel_path: Path) -> str:
|
||||
"""Return the most precise platform tag the wheel is consistent with.
|
||||
|
||||
Mirrors ``auditwheel show`` but returns ``sym_policy`` rather than
|
||||
``overall_policy``: we only care about the glibc symbol versions used,
|
||||
not about other policy axes (ISA extensions, blacklist, etc.) that
|
||||
``overall_policy`` folds in.
|
||||
"""
|
||||
fn = wheel_path.name
|
||||
|
||||
try:
|
||||
arch = get_wheel_architecture(fn)
|
||||
except (WheelToolsError, NonPlatformWheelError):
|
||||
# Architecture isn't deducible from the filename; let auditwheel
|
||||
# infer it from the ELF binaries inside the wheel.
|
||||
arch = None
|
||||
|
||||
try:
|
||||
libc = get_wheel_libc(fn)
|
||||
except WheelToolsError:
|
||||
# An unrepaired wheel uses ``linux_<arch>``, which doesn't encode
|
||||
# libc. Let auditwheel infer it from the ELF binaries.
|
||||
libc = None
|
||||
|
||||
winfo = analyze_wheel_abi(
|
||||
libc,
|
||||
arch,
|
||||
wheel_path,
|
||||
frozenset(),
|
||||
disable_isa_ext_check=False,
|
||||
allow_graft=False,
|
||||
)
|
||||
return winfo.sym_policy.name
|
||||
|
||||
|
||||
def rename_wheel(wheel_path: Path, new_platform_tag: str) -> Path:
|
||||
"""Rename the wheel in place, replacing only its platform tag."""
|
||||
# Wheel filename per PEP 427:
|
||||
# {distribution}-{version}(-{build})?-{python}-{abi}-{platform}.whl
|
||||
# The platform tag is always the last ``-``-separated token before
|
||||
# ``.whl``. Compound tags like ``manylinux_2_31_x86_64`` use ``_`` as the
|
||||
# internal separator, so ``-``-splitting is unambiguous.
|
||||
parts = wheel_path.stem.split("-")
|
||||
if len(parts) < 5:
|
||||
raise ValueError(f"Unrecognised wheel filename: {wheel_path.name}")
|
||||
parts[-1] = new_platform_tag
|
||||
new_path = wheel_path.with_name("-".join(parts) + ".whl")
|
||||
if new_path != wheel_path:
|
||||
wheel_path.rename(new_path)
|
||||
return new_path
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Detect a wheel's manylinux platform tag with "
|
||||
"auditwheel and rename the wheel in place."
|
||||
)
|
||||
parser.add_argument(
|
||||
"wheel",
|
||||
type=Path,
|
||||
help="Path to the wheel to inspect and rename.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
wheel_path: Path = args.wheel
|
||||
if not wheel_path.is_file():
|
||||
print(f"error: {wheel_path} is not a file", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
# Catch the things that ``analyze_wheel_abi`` and ``rename_wheel`` can
|
||||
# raise: any subclass of ``AuditwheelError`` (pure-Python wheels,
|
||||
# invalid libc, malformed wheels), filesystem errors, or our own
|
||||
# ``ValueError`` for an unrecognised wheel filename. Print a single
|
||||
# ``ERROR_TYPE: message`` line to stderr instead of a Python
|
||||
# traceback, which is much friendlier in CI logs.
|
||||
try:
|
||||
new_tag = detect_platform_tag(wheel_path)
|
||||
print(f"detected platform tag: {new_tag}", file=sys.stderr)
|
||||
new_path = rename_wheel(wheel_path, new_tag)
|
||||
except (AuditwheelError, ValueError, OSError) as e:
|
||||
print(
|
||||
f"error: failed to retag {wheel_path.name}: {type(e).__name__}: {e}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 2
|
||||
|
||||
if new_path != wheel_path:
|
||||
print(f"renamed {wheel_path.name} -> {new_path.name}", file=sys.stderr)
|
||||
else:
|
||||
print(f"wheel already tagged {new_tag}", file=sys.stderr)
|
||||
|
||||
print(new_path)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,54 @@
|
||||
#!/bin/bash
|
||||
# Emit docker build flags for release image provenance metadata.
|
||||
# Keep this helper best-effort: missing Buildkite metadata should fall back to
|
||||
# local/default values instead of blocking the Docker build.
|
||||
|
||||
# Variant examples: "", "cu129", "ubuntu2404", "cu129-ubuntu2404".
|
||||
variant="${1:-}"
|
||||
variant_suffix="${variant:+-${variant}}"
|
||||
|
||||
image_name="${VLLM_DOCKER_IMAGE_NAME:-vllm/vllm-openai}"
|
||||
staging_repo="${VLLM_STAGING_IMAGE_REPO:-public.ecr.aws/q9t5s3a7/vllm-release-repo}"
|
||||
build_commit="${VLLM_BUILD_COMMIT:-${BUILDKITE_COMMIT:-unknown}}"
|
||||
build_pipeline="${VLLM_BUILD_PIPELINE:-${BUILDKITE_PIPELINE_ID:-${BUILDKITE_PIPELINE_SLUG:-local}}}"
|
||||
build_url="${VLLM_BUILD_URL:-${BUILDKITE_BUILD_URL:-}}"
|
||||
tag_commit="${BUILDKITE_COMMIT:-${build_commit}}"
|
||||
|
||||
if [[ -n "${BUILDKITE:-}" || -n "${BUILDKITE_COMMIT:-}" ]]; then
|
||||
release_version="${RELEASE_VERSION:-}"
|
||||
if command -v buildkite-agent >/dev/null 2>&1; then
|
||||
release_version="${release_version:-$(buildkite-agent meta-data get release-version 2>/dev/null)}"
|
||||
fi
|
||||
release_version="${release_version#v}"
|
||||
release_version="${release_version:-${tag_commit}}"
|
||||
|
||||
staging_image_ref="${staging_repo}:${tag_commit}-$(uname -m)${variant_suffix}"
|
||||
|
||||
if [[ "${NIGHTLY:-}" == "1" ]]; then
|
||||
if [[ -z "${variant}" ]]; then
|
||||
image_tag="${image_name}:nightly-${tag_commit}"
|
||||
elif [[ "${variant}" == cu* ]]; then
|
||||
cuda_variant="${variant%%-*}"
|
||||
remaining_variant="${variant#${cuda_variant}}"
|
||||
image_tag="${image_name}:${cuda_variant}-nightly-${tag_commit}${remaining_variant}"
|
||||
else
|
||||
image_tag="${image_name}:nightly-${tag_commit}${variant_suffix}"
|
||||
fi
|
||||
else
|
||||
image_tag="${image_name}:v${release_version}${variant_suffix}"
|
||||
fi
|
||||
else
|
||||
image_tag="${VLLM_IMAGE_TAG:-local/vllm-openai:dev}"
|
||||
staging_image_ref="${image_tag}"
|
||||
fi
|
||||
|
||||
emit_arg() {
|
||||
printf -- "--build-arg %s=%s " "$1" "$2"
|
||||
}
|
||||
|
||||
emit_arg VLLM_BUILD_COMMIT "${build_commit}"
|
||||
emit_arg VLLM_BUILD_PIPELINE "${build_pipeline}"
|
||||
emit_arg VLLM_BUILD_URL "${build_url}"
|
||||
# This is the intended public tag. The final digest is only known after push.
|
||||
emit_arg VLLM_IMAGE_TAG "${image_tag}"
|
||||
printf -- "--tag %s " "${staging_image_ref}"
|
||||
@@ -10,20 +10,13 @@ set -ex
|
||||
BUCKET="vllm-wheels"
|
||||
INDICES_OUTPUT_DIR="indices"
|
||||
DEFAULT_VARIANT_ALIAS="cu130" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
|
||||
PYTHON="${PYTHON_PROG:-python3}" # try to read from env var, otherwise use python3
|
||||
SUBPATH=$BUILDKITE_COMMIT
|
||||
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
|
||||
|
||||
# detect if python3.12+ is available
|
||||
has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
|
||||
if [[ "$has_new_python" -eq 0 ]]; then
|
||||
# use new python from docker
|
||||
docker pull python:3-slim
|
||||
PYTHON="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
|
||||
fi
|
||||
|
||||
echo "Using python interpreter: $PYTHON"
|
||||
echo "Python version: $($PYTHON --version)"
|
||||
# Select python3 (>= 3.12) -- local if available, else a docker fallback.
|
||||
# shellcheck source=lib/select-python.sh
|
||||
source .buildkite/scripts/lib/select-python.sh
|
||||
select_python
|
||||
|
||||
# ======== generate and upload indices ========
|
||||
|
||||
|
||||
@@ -3,42 +3,37 @@ set -euox pipefail
|
||||
export VLLM_CPU_CI_ENV=0
|
||||
export VLLM_CPU_KVCACHE_SPACE=1 # avoid OOM
|
||||
|
||||
echo "--- PP+TP"
|
||||
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 --max-model-len=4096 &
|
||||
server_pid=$!
|
||||
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--dataset-name random \
|
||||
--model meta-llama/Llama-3.2-3B-Instruct \
|
||||
--num-prompts 20 \
|
||||
--result-dir ./test_results \
|
||||
--result-filename tp_pp.json \
|
||||
--save-result \
|
||||
--endpoint /v1/completions
|
||||
kill -s SIGTERM $server_pid; wait $server_pid || true
|
||||
failed_req=$(jq '.failed' ./test_results/tp_pp.json)
|
||||
if [ "$failed_req" -ne 0 ]; then
|
||||
echo "Some requests were failed!"
|
||||
exit 1
|
||||
fi
|
||||
MODE=${1:-all}
|
||||
|
||||
echo "--- DP+TP"
|
||||
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
|
||||
server_pid=$!
|
||||
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--dataset-name random \
|
||||
--model meta-llama/Llama-3.2-3B-Instruct \
|
||||
--num-prompts 20 \
|
||||
--result-dir ./test_results \
|
||||
--result-filename dp_pp.json \
|
||||
--save-result \
|
||||
--endpoint /v1/completions
|
||||
kill -s SIGTERM $server_pid; wait $server_pid || true
|
||||
failed_req=$(jq '.failed' ./test_results/dp_pp.json)
|
||||
if [ "$failed_req" -ne 0 ]; then
|
||||
echo "Some requests were failed!"
|
||||
exit 1
|
||||
fi
|
||||
run_scenario() {
|
||||
local label="$1" result_file="$2"
|
||||
shift 2
|
||||
echo "--- $label"
|
||||
vllm serve meta-llama/Llama-3.2-3B-Instruct "$@" --max-model-len=4096 &
|
||||
local server_pid=$!
|
||||
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--dataset-name random \
|
||||
--model meta-llama/Llama-3.2-3B-Instruct \
|
||||
--num-prompts 20 \
|
||||
--result-dir ./test_results \
|
||||
--result-filename "$result_file" \
|
||||
--save-result \
|
||||
--endpoint /v1/completions
|
||||
kill -s SIGTERM "$server_pid"; wait "$server_pid" || true
|
||||
if [ "$(jq '.failed' "./test_results/$result_file")" -ne 0 ]; then
|
||||
echo "Some requests were failed in $label!"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
case "$MODE" in
|
||||
tp_pp) run_scenario "PP+TP" tp_pp.json -tp=2 -pp=2 ;;
|
||||
dp_tp) run_scenario "DP+TP" dp_tp.json -tp=2 -dp=2 ;;
|
||||
all)
|
||||
run_scenario "PP+TP" tp_pp.json -tp=2 -pp=2
|
||||
run_scenario "DP+TP" dp_tp.json -tp=2 -dp=2
|
||||
;;
|
||||
*) echo "ERROR: unknown mode '$MODE' (expected: tp_pp | dp_tp | all)" >&2; exit 1 ;;
|
||||
esac
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Shared helper for rewriting a wheel's platform tag from the generic
|
||||
# ``linux_<arch>`` to the correct ``manylinux_<major>_<minor>_<arch>``.
|
||||
# After sourcing, call ``apply_manylinux_tag <wheel>`` on each wheel
|
||||
# that still carries the generic tag; the renamed path is printed on
|
||||
# stdout (logs go to stderr).
|
||||
#
|
||||
# Why a pinned Docker container instead of using whatever Python
|
||||
# happens to be on the agent:
|
||||
# - vLLM's release agents are heterogeneous -- they don't agree on
|
||||
# a Python minor version, and we can't rely on a particular
|
||||
# ``auditwheel`` being installed.
|
||||
# - ``detect-manylinux-tag.py`` reads ``auditwheel.wheel_abi`` and
|
||||
# ``Policy.sym_policy``, which are *internal* APIs without a
|
||||
# stability promise. Pinning both Python and auditwheel makes the
|
||||
# detected tag a function of the inputs alone, and shifts version
|
||||
# bumps from "implicit drift" to "deliberate, retested change".
|
||||
# - Other release scripts (``generate-and-upload-nightly-index.sh``,
|
||||
# ``upload-rocm-wheels.sh``) already use the python:3-slim image
|
||||
# when the agent's interpreter is too old; this is the same idea
|
||||
# made stricter.
|
||||
#
|
||||
# To keep the per-wheel cost down (the ROCm upload retags ~10 wheels
|
||||
# each run), we install auditwheel into a long-lived helper container
|
||||
# once on source, then ``docker exec`` into it for each call.
|
||||
#
|
||||
# Trap behaviour:
|
||||
# - Sourcing installs an EXIT trap that calls ``manylinux_cleanup`` to
|
||||
# tear down the helper container. Any EXIT trap that was already in
|
||||
# place when this file was sourced is captured and run AFTER our
|
||||
# cleanup, so we don't silently clobber it.
|
||||
# - If a caller sets a new EXIT trap *after* sourcing, that trap will
|
||||
# replace ours; in that case the caller should call
|
||||
# ``manylinux_cleanup`` from their own handler.
|
||||
|
||||
if [[ -n "${_MANYLINUX_LIB_SOURCED:-}" ]]; then
|
||||
return 0
|
||||
fi
|
||||
_MANYLINUX_LIB_SOURCED=1
|
||||
|
||||
# Pin both sides. Bump these deliberately and re-run a representative
|
||||
# wheel from each build target through the detection.
|
||||
_MANYLINUX_PYTHON_IMAGE="python:3.12-slim"
|
||||
_MANYLINUX_AUDITWHEEL_VERSION="6.6.0"
|
||||
|
||||
# Resolve our own directory (and the sibling detect script) using the
|
||||
# canonical, symlink-resolved path. The container mounts cwd at the
|
||||
# same absolute path on both sides, so all paths we hand to it -- the
|
||||
# script, the wheel -- must canonicalise to a location under cwd.
|
||||
_MANYLINUX_LIB_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)"
|
||||
_MANYLINUX_DETECT_SCRIPT="$(cd "${_MANYLINUX_LIB_DIR}/.." && pwd -P)/detect-manylinux-tag.py"
|
||||
_MANYLINUX_CWD="$(pwd -P)"
|
||||
|
||||
docker pull --quiet "$_MANYLINUX_PYTHON_IMAGE" >/dev/null
|
||||
|
||||
# Spin up a long-lived helper container so we install auditwheel once
|
||||
# and then ``docker exec`` into it for each wheel.
|
||||
#
|
||||
# The container runs as root so ``pip install`` can write into the
|
||||
# system site-packages; individual ``docker exec`` calls below pin
|
||||
# themselves to the host UID so any file rename happens with host
|
||||
# ownership, not root.
|
||||
_MANYLINUX_CONTAINER="$(docker run -d --rm \
|
||||
-v "$_MANYLINUX_CWD:$_MANYLINUX_CWD" \
|
||||
-w "$_MANYLINUX_CWD" \
|
||||
"$_MANYLINUX_PYTHON_IMAGE" \
|
||||
sleep infinity)"
|
||||
docker exec "$_MANYLINUX_CONTAINER" \
|
||||
pip install --quiet --disable-pip-version-check \
|
||||
--root-user-action=ignore \
|
||||
"auditwheel==${_MANYLINUX_AUDITWHEEL_VERSION}"
|
||||
|
||||
# Public cleanup -- safe to call multiple times.
|
||||
manylinux_cleanup() {
|
||||
if [[ -n "${_MANYLINUX_CONTAINER:-}" ]]; then
|
||||
docker rm -f "$_MANYLINUX_CONTAINER" >/dev/null 2>&1 || true
|
||||
_MANYLINUX_CONTAINER=""
|
||||
fi
|
||||
}
|
||||
|
||||
# Capture any EXIT trap that was already in place so we can chain to
|
||||
# it rather than overwrite it. ``trap -p EXIT`` prints the handler in
|
||||
# eval-able form (``trap -- 'CMD' EXIT``) or nothing if unset; we
|
||||
# strip the wrapper to recover ``CMD``. Handles the common case --
|
||||
# CMDs without embedded single quotes -- and degrades gracefully (we
|
||||
# still run our own cleanup) for the pathological case.
|
||||
_manylinux_prev_exit_trap_cmd=""
|
||||
_manylinux_existing_exit_trap="$(trap -p EXIT)"
|
||||
if [[ -n "$_manylinux_existing_exit_trap" ]]; then
|
||||
_tmp="${_manylinux_existing_exit_trap#trap -- \'}"
|
||||
_manylinux_prev_exit_trap_cmd="${_tmp%\' EXIT}"
|
||||
unset _tmp
|
||||
fi
|
||||
unset _manylinux_existing_exit_trap
|
||||
|
||||
_manylinux_run_exit_chain() {
|
||||
manylinux_cleanup
|
||||
if [[ -n "$_manylinux_prev_exit_trap_cmd" ]]; then
|
||||
eval "$_manylinux_prev_exit_trap_cmd"
|
||||
fi
|
||||
}
|
||||
trap _manylinux_run_exit_chain EXIT
|
||||
|
||||
# Detect the manylinux platform tag for a single wheel and rename it
|
||||
# in place, printing the renamed wheel path on stdout. Returns
|
||||
# non-zero on failure (which under ``set -e`` propagates to caller).
|
||||
#
|
||||
# The wheel must be reachable via a path under the host cwd so it's
|
||||
# visible inside the helper container; in CI the wheels always live
|
||||
# under ``artifacts/`` so this is fine.
|
||||
apply_manylinux_tag() {
|
||||
local wheel="$1"
|
||||
local abs_wheel
|
||||
abs_wheel="$(realpath "$wheel")"
|
||||
local new_wheel
|
||||
new_wheel="$(docker exec -u "$(id -u):$(id -g)" \
|
||||
"$_MANYLINUX_CONTAINER" \
|
||||
python "$_MANYLINUX_DETECT_SCRIPT" "$abs_wheel")"
|
||||
if [[ -z "$new_wheel" || ! -f "$new_wheel" ]]; then
|
||||
echo "apply_manylinux_tag: detect-manylinux-tag.py did not produce a valid wheel path for $wheel" >&2
|
||||
return 1
|
||||
fi
|
||||
printf '%s\n' "$new_wheel"
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
#!/usr/bin/env bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Pick a Python interpreter for buildkite scripts: prefer a local
|
||||
# ``python3`` if it is recent enough (>= 3.12), otherwise fall back to
|
||||
# a one-shot Docker container running ``python:3-slim``. After
|
||||
# ``select_python`` returns, ``$PYTHON`` is set in the caller's shell
|
||||
# and is safe to use as a command (e.g. ``$PYTHON some_script.py``).
|
||||
#
|
||||
# The 3.12 threshold matches what the existing nightly-index work
|
||||
# expects -- typing features used by ``generate-nightly-index.py``.
|
||||
# This helper does not pin the *minor* version; if you need stricter
|
||||
# reproducibility (e.g. relying on auditwheel internals), invoke
|
||||
# Docker yourself with a pinned tag rather than calling this.
|
||||
|
||||
if [[ -n "${_SELECT_PYTHON_LIB_SOURCED:-}" ]]; then
|
||||
return 0
|
||||
fi
|
||||
_SELECT_PYTHON_LIB_SOURCED=1
|
||||
|
||||
# Sets ``PYTHON`` in the caller's shell and exports it. Idempotent --
|
||||
# calling twice is safe and the second call simply re-runs the probe.
|
||||
select_python() {
|
||||
local py="${PYTHON_PROG:-python3}"
|
||||
local has_new_python
|
||||
has_new_python=$("$py" -c \
|
||||
"print(1 if __import__('sys').version_info >= (3,12) else 0)" \
|
||||
2>/dev/null || echo 0)
|
||||
if [[ "$has_new_python" -eq 0 ]]; then
|
||||
# ``-u $(id -u):$(id -g)`` so files created via the container
|
||||
# end up owned by the host user, not root.
|
||||
docker pull python:3-slim
|
||||
PYTHON="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
|
||||
else
|
||||
PYTHON="$py"
|
||||
fi
|
||||
export PYTHON
|
||||
echo "Using python interpreter: $PYTHON"
|
||||
echo "Python version: $($PYTHON --version)"
|
||||
}
|
||||
@@ -28,6 +28,7 @@
|
||||
# BFCL_MAX_MODEL_LEN - Max model length (default: 4096)
|
||||
# BFCL_PORT - Server port (default: 8000)
|
||||
# BFCL_REASONING_PARSER - Reasoning parser name (default: disabled)
|
||||
# BFCL_TEMPERATURE - Temperature (default: 0.0)
|
||||
# BFCL_EXTRA_ARGS - Additional vLLM server args
|
||||
|
||||
set -euo pipefail
|
||||
@@ -43,6 +44,7 @@ TP_SIZE="${BFCL_TP_SIZE:-1}"
|
||||
MAX_MODEL_LEN="${BFCL_MAX_MODEL_LEN:-4096}"
|
||||
PORT="${BFCL_PORT:-8000}"
|
||||
REASONING_PARSER="${BFCL_REASONING_PARSER:-}"
|
||||
TEMPERATURE="${BFCL_TEMPERATURE:-0.0}"
|
||||
EXTRA_ARGS="${BFCL_EXTRA_ARGS:-}"
|
||||
|
||||
# Set up output directory
|
||||
@@ -139,7 +141,7 @@ echo "vLLM server is ready. (started in ${SECONDS_WAITED}s)"
|
||||
# be patched in-process so BFCL knows to use the OpenAI-compatible handler
|
||||
# against our local vLLM server.
|
||||
bfcl_exit_code=0
|
||||
python3 - "$MODEL" "$TEST_CATEGORY" "$NUM_THREADS" "$PORT" "$API_TYPE" "$OUTPUT_DIR" << 'PYEOF' || bfcl_exit_code=$?
|
||||
python3 - "$MODEL" "$TEST_CATEGORY" "$NUM_THREADS" "$PORT" "$API_TYPE" "$TEMPERATURE" "$OUTPUT_DIR" << 'PYEOF' || bfcl_exit_code=$?
|
||||
import os
|
||||
import sys
|
||||
|
||||
@@ -148,7 +150,8 @@ test_category = sys.argv[2]
|
||||
num_threads = int(sys.argv[3])
|
||||
port = sys.argv[4]
|
||||
api_type = sys.argv[5]
|
||||
output_dir = sys.argv[6] if len(sys.argv) > 6 and sys.argv[6] else os.getcwd()
|
||||
temperature = float(sys.argv[6])
|
||||
output_dir = sys.argv[7] if len(sys.argv) > 7 and sys.argv[7] else os.getcwd()
|
||||
|
||||
os.environ["OPENAI_BASE_URL"] = f"http://localhost:{port}/v1"
|
||||
os.environ["OPENAI_API_KEY"] = "dummy"
|
||||
@@ -204,6 +207,7 @@ gen_kwargs["model"] = [model]
|
||||
gen_kwargs["test_category"] = [c.strip() for c in test_category.split(",")]
|
||||
gen_kwargs["skip_server_setup"] = True
|
||||
gen_kwargs["num_threads"] = num_threads
|
||||
gen_kwargs["temperature"] = temperature
|
||||
generate(**gen_kwargs)
|
||||
|
||||
# ---- evaluate ----
|
||||
|
||||
@@ -2,14 +2,18 @@
|
||||
|
||||
set -ex
|
||||
|
||||
# Upload a single wheel to S3 (rename linux -> manylinux).
|
||||
# Upload a single wheel to S3, after detecting and applying the appropriate
|
||||
# manylinux platform tag with auditwheel.
|
||||
# Index generation is handled separately by generate-and-upload-nightly-index.sh.
|
||||
|
||||
# shellcheck source=lib/manylinux.sh
|
||||
source .buildkite/scripts/lib/manylinux.sh
|
||||
|
||||
BUCKET="vllm-wheels"
|
||||
SUBPATH=$BUILDKITE_COMMIT
|
||||
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
|
||||
|
||||
# ========= collect, rename & upload the wheel ==========
|
||||
# ========= locate the wheel ==========
|
||||
|
||||
# Assume wheels are in artifacts/dist/*.whl
|
||||
wheel_files=(artifacts/dist/*.whl)
|
||||
@@ -21,19 +25,9 @@ if [[ ${#wheel_files[@]} -ne 1 ]]; then
|
||||
fi
|
||||
wheel="${wheel_files[0]}"
|
||||
|
||||
# default build image uses ubuntu 20.04, which corresponds to manylinux_2_31
|
||||
# we also accept params as manylinux tag
|
||||
# refer to https://github.com/mayeut/pep600_compliance?tab=readme-ov-file#acceptable-distros-to-build-wheels
|
||||
manylinux_version="${1:-manylinux_2_31}"
|
||||
# ========= detect manylinux tag and rename ==========
|
||||
|
||||
# Rename 'linux' to the appropriate manylinux version in the wheel filename
|
||||
if [[ "$wheel" != *"linux"* ]]; then
|
||||
echo "Error: Wheel filename does not contain 'linux': $wheel"
|
||||
exit 1
|
||||
fi
|
||||
new_wheel="${wheel/linux/$manylinux_version}"
|
||||
mv -- "$wheel" "$new_wheel"
|
||||
wheel="$new_wheel"
|
||||
wheel="$(apply_manylinux_tag "$wheel")"
|
||||
echo "Renamed wheel to: $wheel"
|
||||
|
||||
# Extract the version from the wheel
|
||||
|
||||
@@ -20,10 +20,6 @@ BUCKET="${S3_BUCKET:-vllm-wheels}"
|
||||
ROCM_SUBPATH="rocm/${BUILDKITE_COMMIT}"
|
||||
S3_COMMIT_PREFIX="s3://$BUCKET/$ROCM_SUBPATH/"
|
||||
INDICES_OUTPUT_DIR="rocm-indices"
|
||||
PYTHON="${PYTHON_PROG:-python3}"
|
||||
|
||||
# ROCm uses manylinux_2_35 (Ubuntu 22.04 based)
|
||||
MANYLINUX_VERSION="manylinux_2_35"
|
||||
|
||||
echo "========================================"
|
||||
echo "ROCm Wheel Upload Configuration"
|
||||
@@ -34,19 +30,21 @@ echo "Commit: $BUILDKITE_COMMIT"
|
||||
echo "Branch: $BUILDKITE_BRANCH"
|
||||
echo "========================================"
|
||||
|
||||
# ======== Part 0: Setup Python ========
|
||||
# ======== Part 0: Setup Python and helpers ========
|
||||
|
||||
# Detect if python3.12+ is available
|
||||
has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)" 2>/dev/null || echo 0)
|
||||
if [[ "$has_new_python" -eq 0 ]]; then
|
||||
# Use new python from docker
|
||||
# Use --user to ensure files are created with correct ownership (not root)
|
||||
docker pull python:3-slim
|
||||
PYTHON="docker run --rm --user $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
|
||||
fi
|
||||
# Pick a Python interpreter for index generation -- local if recent
|
||||
# enough, else a one-shot docker fallback.
|
||||
# shellcheck source=lib/select-python.sh
|
||||
source .buildkite/scripts/lib/select-python.sh
|
||||
select_python
|
||||
|
||||
echo "Using python interpreter: $PYTHON"
|
||||
echo "Python version: $($PYTHON --version)"
|
||||
# Set up auditwheel-in-a-container for the manylinux retagging step.
|
||||
# Distinct from select_python: ``manylinux.sh`` deliberately pins both
|
||||
# the Python and auditwheel versions (the script reads auditwheel
|
||||
# internals) and so always runs in a known-good container regardless
|
||||
# of what's on the agent.
|
||||
# shellcheck source=lib/manylinux.sh
|
||||
source .buildkite/scripts/lib/manylinux.sh
|
||||
|
||||
# ======== Part 1: Collect and prepare wheels ========
|
||||
|
||||
@@ -63,11 +61,18 @@ if [ "$WHEEL_COUNT" -eq 0 ]; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Rename linux to manylinux in wheel filenames
|
||||
# Detect the appropriate manylinux platform tag for any wheel that still
|
||||
# carries the generic ``linux_<arch>`` tag, and rename it in place. We use
|
||||
# auditwheel via ``apply_manylinux_tag`` (see lib/manylinux.sh) rather than
|
||||
# a hard-coded ``manylinux_2_35`` string so that the label tracks the actual
|
||||
# glibc symbol versions used by the binaries (and stays correct if the
|
||||
# rocm_base image is rebased).
|
||||
#
|
||||
# The ``linux``/``manylinux`` filter below skips both pre-tagged wheels
|
||||
# (e.g. upstream torch) and pure-Python ``-any.whl`` wheels.
|
||||
for wheel in all-rocm-wheels/*.whl; do
|
||||
if [[ "$wheel" == *"linux"* ]] && [[ "$wheel" != *"manylinux"* ]]; then
|
||||
new_wheel="${wheel/linux/$MANYLINUX_VERSION}"
|
||||
mv -- "$wheel" "$new_wheel"
|
||||
new_wheel="$(apply_manylinux_tag "$wheel")"
|
||||
echo "Renamed: $(basename "$wheel") -> $(basename "$new_wheel")"
|
||||
fi
|
||||
done
|
||||
|
||||
+17
-17
@@ -395,11 +395,11 @@ steps:
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
- python3 offline_inference/prefix_caching.py
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 offline_inference/llm_engine_example.py
|
||||
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
|
||||
#---------------------------------------------------------- mi250 · kernels ----------------------------------------------------------#
|
||||
|
||||
@@ -1168,13 +1168,13 @@ steps:
|
||||
- vllm/v1/attention/backends/
|
||||
- vllm/v1/attention/selector.py
|
||||
- tests/distributed/test_context_parallel.py
|
||||
- examples/offline_inference/data_parallel.py
|
||||
- examples/features/data_parallel/data_parallel_offline.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=1
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
|
||||
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=allgather_reducescatter --disable-nccl-for-dp-synchronization
|
||||
|
||||
- label: Distributed Tests (4xA100-4xMI300) # TBD
|
||||
timeout_in_minutes: 180
|
||||
@@ -1203,7 +1203,7 @@ steps:
|
||||
- tests/distributed/test_torchrun_example.py
|
||||
- tests/distributed/test_torchrun_example_moe.py
|
||||
- examples/rl/
|
||||
- tests/examples/offline_inference/data_parallel.py
|
||||
- tests/examples/features/data_parallel/data_parallel_offline.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=1
|
||||
@@ -1213,7 +1213,7 @@ steps:
|
||||
- PP_SIZE=2 TP_SIZE=2 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
|
||||
- DP_SIZE=4 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
|
||||
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 distributed/test_torchrun_example_moe.py
|
||||
- python3 ../examples/offline_inference/data_parallel.py --enforce-eager
|
||||
- python3 ../examples/features/data_parallel/data_parallel_offline.py --enforce-eager
|
||||
# rlhf examples
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_nccl.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 ../examples/rl/rlhf_ipc.py
|
||||
@@ -1266,7 +1266,7 @@ steps:
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- examples/offline_inference/torchrun_dp_example.py
|
||||
- examples/features/torchrun/torchrun_dp_example_offline.py
|
||||
- vllm/config/parallel.py
|
||||
- vllm/distributed/
|
||||
- vllm/v1/engine/llm_engine.py
|
||||
@@ -1275,7 +1275,7 @@ steps:
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=1
|
||||
- torchrun --nproc-per-node=8 ../examples/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
|
||||
- torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
|
||||
|
||||
#-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------#
|
||||
|
||||
@@ -1654,11 +1654,11 @@ steps:
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
- python3 offline_inference/prefix_caching.py
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 offline_inference/llm_engine_example.py
|
||||
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
|
||||
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
|
||||
|
||||
@@ -2302,7 +2302,7 @@ steps:
|
||||
commands:
|
||||
- export TORCH_NCCL_BLOCKING_WAIT=1
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
|
||||
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
|
||||
- pytest -v -s tests/distributed/test_packed_tensor.py
|
||||
@@ -2713,7 +2713,7 @@ steps:
|
||||
- vllm/v1/attention/selector.py
|
||||
- tests/distributed/test_context_parallel.py
|
||||
- tests/v1/distributed/test_dbo.py
|
||||
- examples/offline_inference/data_parallel.py
|
||||
- examples/features/data_parallel/data_parallel_offline.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
@@ -2937,11 +2937,11 @@ steps:
|
||||
# Pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# Features demo
|
||||
- python3 offline_inference/prefix_caching.py
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 offline_inference/llm_engine_example.py
|
||||
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
|
||||
#---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------#
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: V1 attention (H100)
|
||||
key: v1-attention-h100
|
||||
timeout_in_minutes: 30
|
||||
device: h100
|
||||
source_file_dependencies:
|
||||
@@ -14,6 +15,7 @@ steps:
|
||||
- pytest -v -s v1/attention
|
||||
|
||||
- label: V1 attention (B200)
|
||||
key: v1-attention-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Basic Correctness
|
||||
key: basic-correctness
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Benchmarks CLI Test
|
||||
key: benchmarks-cli-test
|
||||
timeout_in_minutes: 20
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -12,6 +13,7 @@ steps:
|
||||
- pytest -v -s benchmarks/
|
||||
|
||||
- label: Attention Benchmarks Smoke Test (B200)
|
||||
key: attention-benchmarks-smoke-test-b200
|
||||
device: b200
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Sequence Parallel Correctness Tests (2 GPUs)
|
||||
key: sequence-parallel-correctness-tests-2-gpus
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/"
|
||||
num_devices: 2
|
||||
@@ -17,6 +18,7 @@ steps:
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
|
||||
- label: Sequence Parallel Correctness Tests (2xH100)
|
||||
key: sequence-parallel-correctness-tests-2xh100
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
@@ -27,6 +29,7 @@ steps:
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
|
||||
|
||||
- label: AsyncTP Correctness Tests (2xH100)
|
||||
key: asynctp-correctness-tests-2xh100
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
@@ -37,6 +40,7 @@ steps:
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
|
||||
|
||||
- label: AsyncTP Correctness Tests (B200)
|
||||
key: asynctp-correctness-tests-b200
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
@@ -47,6 +51,7 @@ steps:
|
||||
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
|
||||
|
||||
- label: Distributed Compile Unit Tests (2xH100)
|
||||
key: distributed-compile-unit-tests-2xh100
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
@@ -60,6 +65,7 @@ steps:
|
||||
- pytest -s -v tests/compile/passes/distributed
|
||||
|
||||
- label: Fusion and Compile Unit Tests (2xB200)
|
||||
key: fusion-and-compile-unit-tests-2xb200
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
@@ -89,6 +95,7 @@ steps:
|
||||
- pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
|
||||
|
||||
- label: Fusion E2E Quick (H100)
|
||||
key: fusion-e2e-quick-h100
|
||||
timeout_in_minutes: 15
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
@@ -107,6 +114,7 @@ steps:
|
||||
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and (qwen3 or deepseek)"
|
||||
|
||||
- label: Fusion E2E Config Sweep (H100)
|
||||
key: fusion-e2e-config-sweep-h100
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
@@ -126,6 +134,7 @@ steps:
|
||||
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3"
|
||||
|
||||
- label: Fusion E2E Config Sweep (B200)
|
||||
key: fusion-e2e-config-sweep-b200
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
@@ -139,6 +148,7 @@ steps:
|
||||
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek)) or llama-3)"
|
||||
|
||||
- label: Fusion E2E TP2 Quick (H100)
|
||||
key: fusion-e2e-tp2-quick-h100
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
@@ -156,6 +166,7 @@ steps:
|
||||
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
|
||||
|
||||
- label: Fusion E2E TP2 AR-RMS Config Sweep (H100)
|
||||
key: fusion-e2e-tp2-ar-rms-config-sweep-h100
|
||||
timeout_in_minutes: 40
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
@@ -175,6 +186,7 @@ steps:
|
||||
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3"
|
||||
|
||||
- label: Fusion E2E TP2 AsyncTP Config Sweep (H100)
|
||||
key: fusion-e2e-tp2-asynctp-config-sweep-h100
|
||||
timeout_in_minutes: 40
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: h100
|
||||
@@ -194,6 +206,7 @@ steps:
|
||||
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3"
|
||||
|
||||
- label: Fusion E2E TP2 (B200)
|
||||
key: fusion-e2e-tp2-b200
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Platform Tests (CUDA)
|
||||
key: platform-tests-cuda
|
||||
timeout_in_minutes: 15
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -13,6 +14,7 @@ steps:
|
||||
- pytest -v -s cuda/test_platform_no_cuda_init.py
|
||||
|
||||
- label: Cudagraph
|
||||
key: cudagraph
|
||||
timeout_in_minutes: 20
|
||||
source_file_dependencies:
|
||||
- tests/v1/cudagraph
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Distributed NixlConnector PD accuracy (4 GPUs)
|
||||
key: distributed-nixlconnector-pd-accuracy-4-gpus
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
@@ -13,6 +14,7 @@ steps:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
- label: Distributed FlashInfer NixlConnector PD accuracy (4 GPUs)
|
||||
key: distributed-flashinfer-nixlconnector-pd-accuracy-4-gpus
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
@@ -24,6 +26,7 @@ steps:
|
||||
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
key: dp-ep-distributed-nixlconnector-pd-accuracy-tests-4-gpus
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
@@ -35,6 +38,7 @@ steps:
|
||||
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
key: crosslayer-kv-layout-distributed-nixlconnector-pd-accuracy-tests-4-gpus
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
@@ -46,6 +50,7 @@ steps:
|
||||
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
|
||||
key: hybrid-ssm-nixlconnector-pd-accuracy-tests-4-gpus
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
@@ -57,6 +62,7 @@ steps:
|
||||
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
|
||||
key: multiconnector-nixl-offloading-pd-accuracy-2-gpus
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
@@ -71,6 +77,7 @@ steps:
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
|
||||
|
||||
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
|
||||
key: nixlconnector-pd-spec-decode-acceptance-2-gpus
|
||||
timeout_in_minutes: 30
|
||||
device: a100
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -84,6 +91,7 @@ steps:
|
||||
- bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
|
||||
key: multiconnector-nixl-offloading-pd-edge-cases-2-gpus
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Distributed Comm Ops
|
||||
key: distributed-comm-ops
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
@@ -16,6 +17,7 @@ steps:
|
||||
- pytest -v -s distributed/test_shm_storage.py
|
||||
|
||||
- label: Distributed DP Tests (2 GPUs)
|
||||
key: distributed-dp-tests-2-gpus
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
@@ -37,6 +39,7 @@ steps:
|
||||
- DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Distributed Compile + RPC Tests (2 GPUs)
|
||||
key: distributed-compile-rpc-tests-2-gpus
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
@@ -59,6 +62,7 @@ steps:
|
||||
- pytest -v -s ./compile/test_wrapper.py
|
||||
|
||||
- label: Distributed Torchrun + Shutdown Tests (2 GPUs)
|
||||
key: distributed-torchrun-shutdown-tests-2-gpus
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
@@ -81,6 +85,7 @@ steps:
|
||||
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
|
||||
|
||||
- label: Distributed Torchrun + Examples (4 GPUs)
|
||||
key: distributed-torchrun-examples-4-gpus
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace"
|
||||
num_devices: 4
|
||||
@@ -88,9 +93,8 @@ steps:
|
||||
- vllm/distributed/
|
||||
- tests/distributed/test_torchrun_example.py
|
||||
- tests/distributed/test_torchrun_example_moe.py
|
||||
- examples/offline_inference/rlhf_colocate.py
|
||||
- examples/rl/
|
||||
- tests/examples/offline_inference/data_parallel.py
|
||||
- tests/examples/features/data_parallel/data_parallel_offline.py
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
@@ -107,12 +111,13 @@ steps:
|
||||
# test with torchrun tp=2 and dp=2 with ep
|
||||
- TP_SIZE=2 DP_SIZE=2 ENABLE_EP=1 torchrun --nproc-per-node=4 tests/distributed/test_torchrun_example_moe.py
|
||||
# test with internal dp
|
||||
- python3 examples/offline_inference/data_parallel.py --enforce-eager
|
||||
- python3 examples/features/data_parallel/data_parallel_offline.py --enforce-eager
|
||||
# rlhf examples
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
|
||||
|
||||
- label: Distributed DP Tests (4 GPUs)
|
||||
key: distributed-dp-tests-4-gpus
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
@@ -133,6 +138,7 @@ steps:
|
||||
- pytest -v -s distributed/test_utils.py
|
||||
|
||||
- label: Distributed Compile + Comm (4 GPUs)
|
||||
key: distributed-compile-comm-4-gpus
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
@@ -154,24 +160,28 @@ steps:
|
||||
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
|
||||
|
||||
- label: Distributed Tests (8 GPUs)(H100)
|
||||
key: distributed-tests-8-gpus-h100
|
||||
timeout_in_minutes: 10
|
||||
device: h100
|
||||
num_devices: 8
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- examples/offline_inference/torchrun_dp_example.py
|
||||
- examples/features/torchrun/torchrun_dp_example_offline.py
|
||||
- vllm/config/parallel.py
|
||||
- vllm/distributed/
|
||||
- vllm/v1/engine/llm_engine.py
|
||||
- vllm/v1/executor/uniproc_executor.py
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- tests/distributed/test_mnnvl_alltoall.py
|
||||
|
||||
commands:
|
||||
# https://github.com/NVIDIA/nccl/issues/1838
|
||||
- export NCCL_CUMEM_HOST_ENABLE=0
|
||||
# test with torchrun tp=2 and dp=4 with ep
|
||||
- torchrun --nproc-per-node=8 ../examples/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
|
||||
- torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
|
||||
|
||||
- label: Distributed Tests (4 GPUs)(A100)
|
||||
key: distributed-tests-4-gpus-a100
|
||||
device: a100
|
||||
optional: true
|
||||
num_devices: 4
|
||||
@@ -186,6 +196,7 @@ steps:
|
||||
- pytest -v -s -x lora/test_mixtral.py
|
||||
|
||||
- label: Distributed Tests (2 GPUs)(H100)
|
||||
key: distributed-tests-2-gpus-h100
|
||||
timeout_in_minutes: 15
|
||||
device: h100
|
||||
optional: true
|
||||
@@ -194,12 +205,13 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
|
||||
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/features/data_parallel/data_parallel_offline.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
|
||||
- pytest -v -s tests/distributed/test_packed_tensor.py
|
||||
|
||||
- label: Distributed Tests (2 GPUs)(B200)
|
||||
key: distributed-tests-2-gpus-b200
|
||||
device: b200
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/"
|
||||
@@ -208,8 +220,12 @@ steps:
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
- pytest -v -s tests/distributed/test_mnnvl_alltoall.py
|
||||
|
||||
|
||||
|
||||
- label: 2 Node Test (4 GPUs)
|
||||
key: 2-node-test-4-gpus
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
@@ -222,11 +238,12 @@ steps:
|
||||
- vllm/executor/
|
||||
- vllm/model_executor/models/
|
||||
- tests/distributed/
|
||||
- tests/examples/offline_inference/data_parallel.py
|
||||
- tests/examples/features/data_parallel/data_parallel_offline.py
|
||||
commands:
|
||||
- ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/offline_inference/data_parallel.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
|
||||
- ./.buildkite/scripts/run-multi-node-test.sh /vllm-workspace/tests 2 2 $IMAGE_TAG "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=0 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py && VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py" "VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed' && NUM_NODES=2 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_node_count.py | grep 'Node count test passed' && python3 ../examples/features/data_parallel/data_parallel_offline.py -dp=2 -tp=1 --dp-num-nodes=2 --dp-node-rank=1 --dp-master-addr=192.168.10.10 --dp-master-port=12345 --enforce-eager --trust-remote-code"
|
||||
|
||||
- label: Pipeline + Context Parallelism (4 GPUs)
|
||||
key: pipeline-context-parallelism-4-gpus
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
@@ -241,6 +258,7 @@ steps:
|
||||
- pytest -v -s distributed/test_pipeline_parallel.py
|
||||
|
||||
- label: RayExecutorV2 (4 GPUs)
|
||||
key: rayexecutorv2-4-gpus
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
group: Docker
|
||||
depends_on:
|
||||
- image-build-cpu
|
||||
steps:
|
||||
- label: Docker Build Metadata
|
||||
timeout_in_minutes: 10
|
||||
device: cpu-small
|
||||
source_file_dependencies:
|
||||
- .buildkite/release-pipeline.yaml
|
||||
- .buildkite/scripts/docker-build-metadata-args.sh
|
||||
- docker/Dockerfile
|
||||
- docker/Dockerfile.cpu
|
||||
- docker/docker-bake.hcl
|
||||
- tests/tools/test_docker_build_metadata_args.py
|
||||
commands:
|
||||
- pytest -v -s tools/test_docker_build_metadata_args.py
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: DeepSeek V2-Lite Accuracy
|
||||
key: deepseek-v2-lite-accuracy
|
||||
timeout_in_minutes: 60
|
||||
device: h100
|
||||
optional: true
|
||||
@@ -12,6 +13,7 @@ steps:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010
|
||||
|
||||
- label: Qwen3-30B-A3B-FP8-block Accuracy
|
||||
key: qwen3-30b-a3b-fp8-block-accuracy
|
||||
timeout_in_minutes: 60
|
||||
device: h100
|
||||
optional: true
|
||||
@@ -21,6 +23,7 @@ steps:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
|
||||
|
||||
- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
|
||||
key: qwen3-30b-a3b-fp8-block-accuracy-b200
|
||||
timeout_in_minutes: 60
|
||||
device: b200
|
||||
optional: true
|
||||
@@ -30,6 +33,7 @@ steps:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
|
||||
|
||||
- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy
|
||||
key: qwen3-30b-a3b-fp8-dp4-async-eplb-accuracy
|
||||
timeout_in_minutes: 60
|
||||
device: h100
|
||||
optional: true
|
||||
@@ -39,6 +43,7 @@ steps:
|
||||
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050
|
||||
|
||||
- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
|
||||
key: deepseek-v2-lite-prefetch-offload-accuracy-h100
|
||||
timeout_in_minutes: 60
|
||||
device: h100
|
||||
optional: true
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Engine
|
||||
key: engine
|
||||
timeout_in_minutes: 15
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -16,6 +17,7 @@ steps:
|
||||
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
|
||||
|
||||
- label: Engine (1 GPU)
|
||||
key: engine-1-gpu
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
- vllm/v1/engine/
|
||||
@@ -25,6 +27,7 @@ steps:
|
||||
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
|
||||
|
||||
- label: e2e Scheduling (1 GPU)
|
||||
key: e2e-scheduling-1-gpu
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -34,6 +37,7 @@ steps:
|
||||
- pytest -v -s v1/e2e/general/test_async_scheduling.py
|
||||
|
||||
- label: e2e Core (1 GPU)
|
||||
key: e2e-core-1-gpu
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
- vllm/v1/
|
||||
@@ -42,6 +46,7 @@ steps:
|
||||
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
|
||||
|
||||
- label: V1 e2e (2 GPUs)
|
||||
key: v1-e2e-2-gpus
|
||||
timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
|
||||
optional: true
|
||||
num_devices: 2
|
||||
@@ -58,6 +63,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 e2e (4 GPUs)
|
||||
key: v1-e2e-4-gpus
|
||||
timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
|
||||
optional: true
|
||||
num_devices: 4
|
||||
@@ -74,6 +80,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 e2e (4xH100)
|
||||
key: v1-e2e-4xh100
|
||||
timeout_in_minutes: 60
|
||||
device: h100
|
||||
num_devices: 4
|
||||
|
||||
@@ -2,7 +2,8 @@ group: Entrypoints
|
||||
depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Entrypoints Unit Tests
|
||||
- label: Entrypoints Unit Tests
|
||||
key: entrypoints-unit-tests
|
||||
timeout_in_minutes: 10
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -13,6 +14,7 @@ steps:
|
||||
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling
|
||||
|
||||
- label: Entrypoints Integration (LLM)
|
||||
key: entrypoints-integration-llm
|
||||
timeout_in_minutes: 40
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -26,6 +28,7 @@ steps:
|
||||
- pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 1)
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -43,6 +46,7 @@ steps:
|
||||
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 2)
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -60,6 +64,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server openai - Part 3)
|
||||
key: entrypoints-integration-api-server-openai-part-3
|
||||
timeout_in_minutes: 50
|
||||
device: h200_18gb
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -72,6 +77,7 @@ steps:
|
||||
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
|
||||
|
||||
- label: Entrypoints Integration (API Server 2)
|
||||
key: entrypoints-integration-api-server-2
|
||||
timeout_in_minutes: 130
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -86,6 +92,7 @@ steps:
|
||||
- pytest -v -s tool_use
|
||||
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
key: entrypoints-integration-pooling
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -96,6 +103,7 @@ steps:
|
||||
- pytest -v -s entrypoints/pooling
|
||||
|
||||
- label: Entrypoints Integration (Responses API)
|
||||
key: entrypoints-integration-responses-api
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
@@ -105,6 +113,7 @@ steps:
|
||||
- pytest -v -s entrypoints/openai/responses
|
||||
|
||||
- label: OpenAI API Correctness
|
||||
key: openai-api-correctness
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: EPLB Algorithm
|
||||
key: eplb-algorithm
|
||||
timeout_in_minutes: 15
|
||||
device: h200_18gb
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -15,6 +16,7 @@ steps:
|
||||
- pytest -v -s distributed/test_eplb_utils.py
|
||||
|
||||
- label: EPLB Execution # 17min
|
||||
key: eplb-execution
|
||||
timeout_in_minutes: 27
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
@@ -26,6 +28,7 @@ steps:
|
||||
- pytest -v -s distributed/test_eplb_spec_decode.py
|
||||
|
||||
- label: Elastic EP Scaling Test
|
||||
key: elastic-ep-scaling-test
|
||||
timeout_in_minutes: 20
|
||||
device: h100
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: vLLM IR Tests
|
||||
key: vllm-ir-tests
|
||||
timeout_in_minutes: 10
|
||||
device: h200_18gb
|
||||
working_dir: "/vllm-workspace/"
|
||||
@@ -14,6 +15,7 @@ steps:
|
||||
- pytest -v -s tests/kernels/ir
|
||||
|
||||
- label: Kernels Core Operation Test
|
||||
key: kernels-core-operation-test
|
||||
timeout_in_minutes: 75
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
@@ -23,6 +25,7 @@ steps:
|
||||
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py
|
||||
|
||||
- label: Kernels MiniMax Reduce RMS Test (2 GPUs)
|
||||
key: kernels-minimax-reduce-rms-test-2-gpus
|
||||
timeout_in_minutes: 15
|
||||
num_devices: 2
|
||||
device: h100
|
||||
@@ -36,6 +39,7 @@ steps:
|
||||
- pytest -v -s kernels/core/test_minimax_reduce_rms.py
|
||||
|
||||
- label: Kernels Attention Test %N
|
||||
key: kernels-attention-test
|
||||
timeout_in_minutes: 35
|
||||
source_file_dependencies:
|
||||
- csrc/attention/
|
||||
@@ -49,6 +53,7 @@ steps:
|
||||
parallelism: 2
|
||||
|
||||
- label: Kernels Quantization Test %N
|
||||
key: kernels-quantization-test
|
||||
timeout_in_minutes: 90
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/
|
||||
@@ -59,6 +64,7 @@ steps:
|
||||
parallelism: 2
|
||||
|
||||
- label: Kernels MoE Test %N
|
||||
key: kernels-moe-test
|
||||
timeout_in_minutes: 25
|
||||
source_file_dependencies:
|
||||
- csrc/quantization/cutlass_w8a8/moe/
|
||||
@@ -74,6 +80,7 @@ steps:
|
||||
parallelism: 5
|
||||
|
||||
- label: Kernels Mamba Test
|
||||
key: kernels-mamba-test
|
||||
timeout_in_minutes: 45
|
||||
source_file_dependencies:
|
||||
- csrc/mamba/
|
||||
@@ -82,7 +89,18 @@ steps:
|
||||
commands:
|
||||
- pytest -v -s kernels/mamba
|
||||
|
||||
- label: Kernels KDA Test
|
||||
timeout_in_minutes: 20
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/fla/ops/kda.py
|
||||
- vllm/model_executor/layers/fla/ops/chunk_delta_h.py
|
||||
- vllm/model_executor/layers/fla/ops/l2norm.py
|
||||
- tests/kernels/test_kda.py
|
||||
commands:
|
||||
- pytest -v -s kernels/test_kda.py
|
||||
|
||||
- label: Kernels DeepGEMM Test (H100)
|
||||
key: kernels-deepgemm-test-h100
|
||||
timeout_in_minutes: 45
|
||||
device: h100
|
||||
num_devices: 1
|
||||
@@ -104,6 +122,7 @@ steps:
|
||||
- pytest -v -s quantization/test_cutlass_w4a16.py
|
||||
|
||||
- label: Kernels (B200)
|
||||
key: kernels-b200
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
@@ -152,6 +171,7 @@ steps:
|
||||
- pytest -v -s tests/models/quantization/test_nvfp4.py
|
||||
|
||||
- label: Kernels Helion Test
|
||||
key: kernels-helion-test
|
||||
timeout_in_minutes: 30
|
||||
device: h100
|
||||
source_file_dependencies:
|
||||
@@ -163,6 +183,7 @@ steps:
|
||||
|
||||
|
||||
- label: Kernels FP8 MoE Test (1 H100)
|
||||
key: kernels-fp8-moe-test-1-h100
|
||||
timeout_in_minutes: 90
|
||||
device: h100
|
||||
num_devices: 1
|
||||
@@ -179,6 +200,7 @@ steps:
|
||||
- pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py
|
||||
|
||||
- label: Kernels FP8 MoE Test (2 H100s)
|
||||
key: kernels-fp8-moe-test-2-h100s
|
||||
timeout_in_minutes: 90
|
||||
device: h100
|
||||
num_devices: 2
|
||||
@@ -188,6 +210,7 @@ steps:
|
||||
- pytest -v -s kernels/moe/test_deepep_moe.py
|
||||
|
||||
- label: Kernels Fp4 MoE Test (B200)
|
||||
key: kernels-fp4-moe-test-b200
|
||||
timeout_in_minutes: 60
|
||||
device: b200
|
||||
num_devices: 1
|
||||
@@ -200,6 +223,7 @@ steps:
|
||||
|
||||
|
||||
- label: Kernels FusedMoE Layer Test (2 H100s)
|
||||
key: kernels-fusedmoe-layer-test-2-h100s
|
||||
timeout_in_minutes: 90
|
||||
device: h100
|
||||
num_devices: 2
|
||||
@@ -216,6 +240,7 @@ steps:
|
||||
|
||||
|
||||
- label: Kernels FusedMoE Layer Test (2 B200s)
|
||||
key: kernels-fusedmoe-layer-test-2-b200s
|
||||
timeout_in_minutes: 90
|
||||
device: b200
|
||||
num_devices: 2
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: LM Eval Small Models
|
||||
key: lm-eval-small-models
|
||||
timeout_in_minutes: 75
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
@@ -24,6 +25,7 @@ steps:
|
||||
# - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
|
||||
|
||||
- label: LM Eval Large Models (4 GPUs)(H100)
|
||||
key: lm-eval-large-models-4-gpus-h100
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 4
|
||||
@@ -36,6 +38,7 @@ steps:
|
||||
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
|
||||
|
||||
- label: LM Eval Small Models (B200)
|
||||
key: lm-eval-small-models-b200
|
||||
timeout_in_minutes: 120
|
||||
device: b200
|
||||
optional: true
|
||||
@@ -46,6 +49,7 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
|
||||
|
||||
- label: LM Eval Qwen3.5 Models (B200)
|
||||
key: lm-eval-qwen3-5-models-b200
|
||||
timeout_in_minutes: 120
|
||||
device: b200
|
||||
optional: true
|
||||
@@ -62,6 +66,7 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
|
||||
|
||||
- label: LM Eval Large Models (H200)
|
||||
key: lm-eval-large-models-h200
|
||||
timeout_in_minutes: 60
|
||||
device: h200
|
||||
optional: true
|
||||
@@ -70,6 +75,7 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
|
||||
|
||||
- label: MoE Refactor Integration Test (H100 - TEMPORARY)
|
||||
key: moe-refactor-integration-test-h100-temporary
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 2
|
||||
@@ -77,6 +83,7 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100.txt
|
||||
|
||||
- label: MoE Refactor Integration Test (B200 - TEMPORARY)
|
||||
key: moe-refactor-integration-test-b200-temporary
|
||||
device: b200
|
||||
optional: true
|
||||
num_devices: 2
|
||||
@@ -84,6 +91,7 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt
|
||||
|
||||
- label: MoE Refactor Integration Test (B200 DP - TEMPORARY)
|
||||
key: moe-refactor-integration-test-b200-dp-temporary
|
||||
device: b200
|
||||
optional: true
|
||||
num_devices: 2
|
||||
@@ -92,6 +100,7 @@ steps:
|
||||
|
||||
|
||||
- label: LM Eval TurboQuant KV Cache
|
||||
key: lm-eval-turboquant-kv-cache
|
||||
timeout_in_minutes: 75
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization/turboquant/
|
||||
@@ -102,6 +111,7 @@ steps:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
|
||||
|
||||
- label: GPQA Eval (GPT-OSS) (H100)
|
||||
key: gpqa-eval-gpt-oss-h100
|
||||
timeout_in_minutes: 120
|
||||
device: h100
|
||||
optional: true
|
||||
@@ -115,6 +125,7 @@ steps:
|
||||
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt
|
||||
|
||||
- label: GPQA Eval (GPT-OSS) (B200)
|
||||
key: gpqa-eval-gpt-oss-b200
|
||||
timeout_in_minutes: 120
|
||||
device: b200
|
||||
optional: true
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: LoRA %N
|
||||
key: lora
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
- vllm/lora
|
||||
@@ -13,6 +14,7 @@ steps:
|
||||
|
||||
|
||||
- label: LoRA TP (Distributed)
|
||||
key: lora-tp-distributed
|
||||
timeout_in_minutes: 30
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: V1 Spec Decode
|
||||
key: v1-spec-decode
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -18,6 +19,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 Sample + Logits
|
||||
key: v1-sample-logits
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -41,6 +43,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 Core + KV + Metrics
|
||||
key: v1-core-kv-metrics
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -71,6 +74,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 Others (CPU)
|
||||
key: v1-others-cpu
|
||||
depends_on:
|
||||
- image-build-cpu
|
||||
source_file_dependencies:
|
||||
@@ -86,6 +90,7 @@ steps:
|
||||
- pytest -v -s -m 'cpu_test' v1/metrics
|
||||
|
||||
- label: Regression
|
||||
key: regression
|
||||
timeout_in_minutes: 20
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -97,6 +102,7 @@ steps:
|
||||
working_dir: "/vllm-workspace/tests" # optional
|
||||
|
||||
- label: Examples
|
||||
key: examples
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/examples"
|
||||
source_file_dependencies:
|
||||
@@ -120,14 +126,15 @@ steps:
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
- python3 offline_inference/prefix_caching.py
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 offline_inference/llm_engine_example.py
|
||||
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
# https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
|
||||
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
|
||||
- label: Metrics, Tracing (2 GPUs)
|
||||
key: metrics-tracing-2-gpus
|
||||
timeout_in_minutes: 20
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
@@ -142,6 +149,7 @@ steps:
|
||||
- pytest -v -s v1/tracing
|
||||
|
||||
- label: Python-only Installation
|
||||
key: python-only-installation
|
||||
depends_on: ~
|
||||
timeout_in_minutes: 20
|
||||
source_file_dependencies:
|
||||
@@ -151,6 +159,7 @@ steps:
|
||||
- bash standalone_tests/python_only_compile.sh
|
||||
|
||||
- label: Async Engine, Inputs, Utils, Worker
|
||||
key: async-engine-inputs-utils-worker
|
||||
timeout_in_minutes: 50
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -163,7 +172,8 @@ steps:
|
||||
- pytest -v -s utils_
|
||||
|
||||
- label: Async Engine, Inputs, Utils, Worker, Config (CPU)
|
||||
depends_on:
|
||||
key: async-engine-inputs-utils-worker-config-cpu
|
||||
depends_on:
|
||||
- image-build-cpu
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
@@ -196,6 +206,7 @@ steps:
|
||||
- pytest -v -s config
|
||||
|
||||
- label: Batch Invariance (H100)
|
||||
key: batch-invariance-h100
|
||||
timeout_in_minutes: 30
|
||||
device: h100
|
||||
source_file_dependencies:
|
||||
@@ -211,6 +222,7 @@ steps:
|
||||
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
|
||||
|
||||
- label: Batch Invariance (B200)
|
||||
key: batch-invariance-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
source_file_dependencies:
|
||||
@@ -227,6 +239,7 @@ steps:
|
||||
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
|
||||
|
||||
- label: Acceptance Length Test (Large Models) # optional
|
||||
key: acceptance-length-test-large-models
|
||||
timeout_in_minutes: 25
|
||||
gpu: h100
|
||||
optional: true
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Model Executor
|
||||
key: model-executor
|
||||
timeout_in_minutes: 35
|
||||
source_file_dependencies:
|
||||
- vllm/engine/arg_utils.py
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Model Runner V2 Core Tests
|
||||
key: model-runner-v2-core-tests
|
||||
timeout_in_minutes: 45
|
||||
source_file_dependencies:
|
||||
- vllm/v1/worker/gpu/
|
||||
@@ -25,14 +26,16 @@ steps:
|
||||
- pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
|
||||
|
||||
- label: Model Runner V2 Examples
|
||||
key: model-runner-v2-examples
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/examples"
|
||||
source_file_dependencies:
|
||||
- vllm/v1/worker/gpu/
|
||||
- vllm/v1/core/sched/
|
||||
- vllm/v1/worker/gpu_worker.py
|
||||
- examples/offline_inference/
|
||||
- examples/basic/offline_inference/
|
||||
- examples/generate/multimodal/
|
||||
- examples/features/
|
||||
- examples/pooling/embed/vision_embedding_offline.py
|
||||
- examples/others/tensorize_vllm_model.py
|
||||
commands:
|
||||
@@ -51,14 +54,15 @@ steps:
|
||||
# for pooling models
|
||||
- python3 pooling/embed/vision_embedding_offline.py --seed 0
|
||||
# for features demo
|
||||
- python3 offline_inference/prefix_caching.py
|
||||
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
|
||||
- python3 offline_inference/llm_engine_example.py
|
||||
- python3 others/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 others/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
|
||||
- python3 offline_inference/spec_decode.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 2048
|
||||
# https://github.com/vllm-project/vllm/pull/26682 uses slightly more memory in PyTorch 2.9+ causing this test to OOM in 1xL4 GPU
|
||||
- python3 offline_inference/spec_decode.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
- python3 features/speculative_decoding/spec_decode_offline.py --test --method eagle3 --num_spec_tokens 3 --dataset-name hf --dataset-path philschmid/mt-bench --num-prompts 80 --temp 0 --top-p 1.0 --top-k -1 --tp 1 --enable-chunked-prefill --max-model-len 1536
|
||||
|
||||
- label: Model Runner V2 Distributed (2 GPUs)
|
||||
key: model-runner-v2-distributed-2-gpus
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
@@ -79,6 +83,7 @@ steps:
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
|
||||
- label: Model Runner V2 Pipeline Parallelism (4 GPUs)
|
||||
key: model-runner-v2-pipeline-parallelism-4-gpus
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
@@ -94,6 +99,7 @@ steps:
|
||||
- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
|
||||
|
||||
- label: Model Runner V2 Spec Decode
|
||||
key: model-runner-v2-spec-decode
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Basic Models Tests (Initialization)
|
||||
key: basic-models-tests-initialization
|
||||
timeout_in_minutes: 45
|
||||
torch_nightly: true
|
||||
source_file_dependencies:
|
||||
@@ -16,6 +17,7 @@ steps:
|
||||
torch_nightly: {}
|
||||
|
||||
- label: Basic Models Tests (Extra Initialization) %N
|
||||
key: basic-models-tests-extra-initialization
|
||||
timeout_in_minutes: 45
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/models/
|
||||
@@ -31,6 +33,7 @@ steps:
|
||||
torch_nightly: {}
|
||||
|
||||
- label: Basic Models Tests (Other)
|
||||
key: basic-models-tests-other
|
||||
timeout_in_minutes: 45
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -47,6 +50,7 @@ steps:
|
||||
|
||||
|
||||
- label: Basic Models Test (Other CPU) # 5min
|
||||
key: basic-models-test-other-cpu
|
||||
depends_on:
|
||||
- image-build-cpu
|
||||
timeout_in_minutes: 10
|
||||
@@ -59,6 +63,7 @@ steps:
|
||||
- pytest -v -s models/test_utils.py models/test_vision.py
|
||||
|
||||
- label: Transformers Nightly Models
|
||||
key: transformers-nightly-models
|
||||
working_dir: "/vllm-workspace/"
|
||||
optional: true
|
||||
soft_fail: true
|
||||
@@ -74,6 +79,7 @@ steps:
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
- label: Transformers Backward Compatibility Models Test
|
||||
key: transformers-backward-compatibility-models-test
|
||||
working_dir: "/vllm-workspace/"
|
||||
optional: true
|
||||
soft_fail: true
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Distributed Model Tests (2 GPUs)
|
||||
key: distributed-model-tests-2-gpus
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Language Models Tests (Standard)
|
||||
key: language-models-tests-standard
|
||||
timeout_in_minutes: 25
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -15,6 +16,7 @@ steps:
|
||||
torch_nightly: {}
|
||||
|
||||
- label: Language Models Tests (Extra Standard) %N
|
||||
key: language-models-tests-extra-standard
|
||||
timeout_in_minutes: 45
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/models/
|
||||
@@ -31,6 +33,7 @@ steps:
|
||||
torch_nightly: {}
|
||||
|
||||
- label: Language Models Tests (Hybrid) %N
|
||||
key: language-models-tests-hybrid
|
||||
timeout_in_minutes: 75
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -47,6 +50,7 @@ steps:
|
||||
torch_nightly: {}
|
||||
|
||||
- label: Language Models Test (Extended Generation) # 80min
|
||||
key: language-models-test-extended-generation
|
||||
timeout_in_minutes: 110
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
@@ -69,6 +73,7 @@ steps:
|
||||
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
|
||||
|
||||
- label: Language Models Test (PPL)
|
||||
key: language-models-test-ppl
|
||||
timeout_in_minutes: 110
|
||||
device: h200_18gb
|
||||
optional: true
|
||||
@@ -79,6 +84,7 @@ steps:
|
||||
- pytest -v -s models/language/generation_ppl_test
|
||||
|
||||
- label: Language Models Test (Extended Pooling) # 36min
|
||||
key: language-models-test-extended-pooling
|
||||
timeout_in_minutes: 50
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
@@ -93,6 +99,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Language Models Test (MTEB)
|
||||
key: language-models-test-mteb
|
||||
timeout_in_minutes: 110
|
||||
device: h200_18gb
|
||||
optional: true
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: "Multi-Modal Models (Standard) 1: qwen2"
|
||||
key: multi-modal-models-standard-1-qwen2
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -19,6 +20,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
|
||||
key: multi-modal-models-standard-2-qwen3-gemma
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -36,6 +38,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
|
||||
key: multi-modal-models-standard-3-llava-qwen2-vl
|
||||
timeout_in_minutes: 45
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -51,6 +54,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 4: other + whisper"
|
||||
key: multi-modal-models-standard-4-other-whisper
|
||||
timeout_in_minutes: 45
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -67,7 +71,8 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Multi-Modal Processor (CPU)
|
||||
depends_on:
|
||||
key: multi-modal-processor-cpu
|
||||
depends_on:
|
||||
- image-build-cpu
|
||||
timeout_in_minutes: 60
|
||||
source_file_dependencies:
|
||||
@@ -80,6 +85,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
|
||||
|
||||
- label: Multi-Modal Processor # 44min
|
||||
key: multi-modal-processor
|
||||
timeout_in_minutes: 60
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -91,6 +97,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/processing/test_tensor_schema.py
|
||||
|
||||
- label: Multi-Modal Accuracy Eval (Small Models) # 50min
|
||||
key: multi-modal-accuracy-eval-small-models
|
||||
timeout_in_minutes: 70
|
||||
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
|
||||
source_file_dependencies:
|
||||
@@ -101,6 +108,7 @@ steps:
|
||||
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
|
||||
|
||||
- label: Multi-Modal Models (Extended Generation 1)
|
||||
key: multi-modal-models-extended-generation-1
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -117,6 +125,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Multi-Modal Models (Extended Generation 2)
|
||||
key: multi-modal-models-extended-generation-2
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -126,6 +135,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
|
||||
|
||||
- label: Multi-Modal Models (Extended Generation 3)
|
||||
key: multi-modal-models-extended-generation-3
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -135,6 +145,7 @@ steps:
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
|
||||
|
||||
- label: Multi-Modal Models (Extended Pooling)
|
||||
key: multi-modal-models-extended-pooling
|
||||
optional: true
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Plugin Tests (2 GPUs)
|
||||
key: plugin-tests-2-gpus
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: PyTorch Compilation Unit Tests
|
||||
key: pytorch-compilation-unit-tests
|
||||
timeout_in_minutes: 10
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -18,6 +19,7 @@ steps:
|
||||
- "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
|
||||
|
||||
- label: PyTorch Compilation Unit Tests (H100)
|
||||
key: pytorch-compilation-unit-tests-h100
|
||||
timeout_in_minutes: 30
|
||||
device: h100
|
||||
num_devices: 1
|
||||
@@ -28,6 +30,7 @@ steps:
|
||||
- "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
|
||||
|
||||
- label: PyTorch Compilation Passes Unit Tests
|
||||
key: pytorch-compilation-passes-unit-tests
|
||||
timeout_in_minutes: 20
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -36,6 +39,7 @@ steps:
|
||||
- pytest -s -v compile/passes --ignore compile/passes/distributed
|
||||
|
||||
- label: PyTorch Fullgraph Smoke Test
|
||||
key: pytorch-fullgraph-smoke-test
|
||||
timeout_in_minutes: 35
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -48,6 +52,7 @@ steps:
|
||||
- "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
|
||||
|
||||
- label: PyTorch Fullgraph
|
||||
key: pytorch-fullgraph
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -58,6 +63,7 @@ steps:
|
||||
- pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
|
||||
|
||||
- label: Pytorch Nightly Dependency Override Check # 2min
|
||||
key: pytorch-nightly-dependency-override-check
|
||||
# if this test fails, it means the nightly torch version is not compatible with some
|
||||
# of the dependencies. Please check the error message and add the package to whitelist
|
||||
# in /vllm/tools/pre_commit/generate_nightly_torch_test.py
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Quantization
|
||||
key: quantization
|
||||
timeout_in_minutes: 90
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
@@ -21,6 +22,7 @@ steps:
|
||||
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
|
||||
|
||||
- label: Quantized MoE Test (B200)
|
||||
key: quantized-moe-test-b200
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/"
|
||||
device: b200
|
||||
@@ -38,6 +40,7 @@ steps:
|
||||
- pytest -s -v tests/quantization/test_blackwell_moe.py
|
||||
|
||||
- label: Quantized Models Test
|
||||
key: quantized-models-test
|
||||
timeout_in_minutes: 60
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Ray Dependency Compatibility Check
|
||||
key: ray-dependency-compatibility-check
|
||||
# Informational only — does not block the pipeline.
|
||||
# If this fails, it means the PR introduces a dependency that
|
||||
# conflicts with Ray's dependency constraints.
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Samplers Test
|
||||
key: samplers-test
|
||||
timeout_in_minutes: 75
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers
|
||||
@@ -10,7 +11,9 @@ steps:
|
||||
- tests/samplers
|
||||
- tests/conftest.py
|
||||
commands:
|
||||
- pytest -v -s samplers
|
||||
# VLLM_USE_FLASHINFER_SAMPLER defaults to 1 now, so we need to pin both
|
||||
# values explicitly to still cover the PyTorch-native (Triton) path.
|
||||
- VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers
|
||||
- VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
|
||||
mirror:
|
||||
amd:
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Spec Decode Eagle
|
||||
key: spec-decode-eagle
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -13,6 +14,7 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
|
||||
|
||||
- label: Spec Decode Eagle Nightly B200
|
||||
key: spec-decode-eagle-nightly-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
optional: true
|
||||
@@ -24,6 +26,7 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
|
||||
|
||||
- label: Spec Decode Speculators + MTP
|
||||
key: spec-decode-speculators-mtp
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -35,6 +38,7 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
|
||||
|
||||
- label: Spec Decode Speculators + MTP Nightly B200
|
||||
key: spec-decode-speculators-mtp-nightly-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
optional: true
|
||||
@@ -47,6 +51,7 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
|
||||
|
||||
- label: Spec Decode Ngram + Suffix
|
||||
key: spec-decode-ngram-suffix
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -57,6 +62,7 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
|
||||
|
||||
- label: Spec Decode Draft Model
|
||||
key: spec-decode-draft-model
|
||||
timeout_in_minutes: 30
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
@@ -67,6 +73,7 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
|
||||
|
||||
- label: Spec Decode Draft Model Nightly B200
|
||||
key: spec-decode-draft-model-nightly-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
optional: true
|
||||
@@ -78,6 +85,7 @@ steps:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
|
||||
|
||||
- label: DFlash Speculators Correctness
|
||||
key: dflash-speculators-correctness
|
||||
timeout_in_minutes: 30
|
||||
device: h100
|
||||
optional: true
|
||||
@@ -89,3 +97,16 @@ steps:
|
||||
commands:
|
||||
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
|
||||
- pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
|
||||
|
||||
- label: Spec Decode MTP hybrid (B200)
|
||||
timeout_in_minutes: 30
|
||||
device: b200
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/v1/worker/gpu/spec_decode/
|
||||
- vllm/model_executor/models/qwen3_5.py
|
||||
- vllm/model_executor/models/qwen3_5_mtp.py
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Weight Loading Multiple GPU # 33min
|
||||
key: weight-loading-multiple-gpu
|
||||
timeout_in_minutes: 45
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
|
||||
+2
-3
@@ -308,8 +308,7 @@ pull_request_rules:
|
||||
- files=benchmarks/benchmark_serving_structured_output.py
|
||||
- files=benchmarks/run_structured_output_benchmark.sh
|
||||
- files=docs/features/structured_outputs.md
|
||||
- files=examples/offline_inference/structured_outputs.py
|
||||
- files=examples/online_serving/structured_outputs/structured_outputs.py
|
||||
- files=^examples/features/structured_outputs/
|
||||
- files~=^tests/v1/structured_output/
|
||||
- files=tests/entrypoints/llm/test_struct_output_generate.py
|
||||
- files~=^vllm/v1/structured_output/
|
||||
@@ -325,7 +324,7 @@ pull_request_rules:
|
||||
- or:
|
||||
- files~=^vllm/v1/spec_decode/
|
||||
- files~=^tests/v1/spec_decode/
|
||||
- files~=^examples/.*(spec_decode|mlpspeculator|eagle|speculation).*\.py
|
||||
- files=^examples/features/speculative_decoding/
|
||||
- files~=^vllm/model_executor/models/.*eagle.*\.py
|
||||
- files=vllm/model_executor/models/mlp_speculator.py
|
||||
- files~=^vllm/transformers_utils/configs/(eagle|medusa|mlp_speculator)\.py
|
||||
|
||||
@@ -16,11 +16,7 @@ permissions:
|
||||
|
||||
jobs:
|
||||
pre-run-check:
|
||||
if: >-
|
||||
github.event_name == 'pull_request' &&
|
||||
(github.event.action != 'labeled' ||
|
||||
github.event.label.name == 'ready' ||
|
||||
github.event.label.name == 'verified')
|
||||
if: github.event_name == 'pull_request'
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check PR label and author merge count
|
||||
@@ -49,12 +45,7 @@ jobs:
|
||||
|
||||
pre-commit:
|
||||
needs: pre-run-check
|
||||
if: >-
|
||||
always() &&
|
||||
(github.event.action != 'labeled' ||
|
||||
github.event.label.name == 'ready' ||
|
||||
github.event.label.name == 'verified') &&
|
||||
(needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
|
||||
if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
|
||||
|
||||
@@ -237,6 +237,7 @@ ep_kernels_workspace/
|
||||
|
||||
# Allow tracked library source folders under submodules (e.g., benchmarks/lib)
|
||||
!vllm/benchmarks/lib/
|
||||
!.buildkite/scripts/lib/
|
||||
|
||||
# Generated gRPC protobuf files (compiled at build time from vllm_engine.proto)
|
||||
vllm/grpc/vllm_engine_pb2.py
|
||||
|
||||
@@ -217,6 +217,7 @@ async def send_request(
|
||||
min_tokens: int | None = None,
|
||||
max_tokens: int | None = None,
|
||||
timeout_sec: int = 120,
|
||||
conversation_id: str | None = None,
|
||||
) -> ServerResponse:
|
||||
payload = {
|
||||
"model": model,
|
||||
@@ -225,6 +226,9 @@ async def send_request(
|
||||
"temperature": 0.0,
|
||||
}
|
||||
|
||||
if conversation_id is not None:
|
||||
payload["conversation_id"] = conversation_id
|
||||
|
||||
if stream:
|
||||
payload["stream"] = True
|
||||
payload["stream_options"] = {"include_usage": False}
|
||||
@@ -419,6 +423,7 @@ async def send_turn(
|
||||
min_tokens,
|
||||
max_tokens,
|
||||
req_args.timeout_sec,
|
||||
conversation_id=conv_id,
|
||||
)
|
||||
|
||||
if response.valid is False:
|
||||
|
||||
+78
-25
@@ -1,5 +1,16 @@
|
||||
#include "cpu_attn_dispatch_generated.h"
|
||||
|
||||
// Maps kv_cache_dtype string to Fp8KVCacheDataType enum.
|
||||
// "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2.
|
||||
static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
|
||||
const std::string& kv_cache_dtype) {
|
||||
if (kv_cache_dtype == "fp8_e5m2")
|
||||
return cpu_attention::Fp8KVCacheDataType::kFp8E5M2;
|
||||
if (kv_cache_dtype == "fp8_e4m3" || kv_cache_dtype == "fp8")
|
||||
return cpu_attention::Fp8KVCacheDataType::kFp8E4M3;
|
||||
return cpu_attention::Fp8KVCacheDataType::kAuto;
|
||||
}
|
||||
|
||||
torch::Tensor get_scheduler_metadata(
|
||||
const int64_t num_req, const int64_t num_heads_q,
|
||||
const int64_t num_heads_kv, const int64_t head_dim,
|
||||
@@ -49,7 +60,7 @@ torch::Tensor get_scheduler_metadata(
|
||||
input.enable_kv_split = enable_kv_split;
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
|
||||
CPU_ATTN_DISPATCH(head_dim, isa, [&]() {
|
||||
CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
|
||||
input.elem_size = sizeof(scalar_t);
|
||||
input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
|
||||
input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
|
||||
@@ -72,7 +83,9 @@ void cpu_attn_reshape_and_cache(
|
||||
key_cache, // [num_blocks, num_kv_heads, block_size, head_size]
|
||||
torch::Tensor&
|
||||
value_cache, // [num_blocks, num_kv_heads, block_size, head_size]
|
||||
const torch::Tensor& slot_mapping, const std::string& isa) {
|
||||
const torch::Tensor& slot_mapping, const std::string& isa,
|
||||
const double k_scale = 1.0, const double v_scale = 1.0,
|
||||
const std::string& kv_cache_dtype = "auto") {
|
||||
TORCH_CHECK_EQ(key.dim(), 3);
|
||||
TORCH_CHECK_EQ(value.dim(), 3);
|
||||
TORCH_CHECK_EQ(key_cache.dim(), 4);
|
||||
@@ -80,18 +93,30 @@ void cpu_attn_reshape_and_cache(
|
||||
TORCH_CHECK_EQ(key.stride(2), 1);
|
||||
TORCH_CHECK_EQ(value.stride(2), 1);
|
||||
|
||||
const int64_t kv_cache_idx =
|
||||
static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
|
||||
const bool is_fp8 = (kv_cache_idx != 0);
|
||||
|
||||
if (is_fp8) {
|
||||
TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
|
||||
"key_cache must be uint8 for FP8 path");
|
||||
TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
|
||||
"value_cache must be uint8 for FP8 path");
|
||||
TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
|
||||
TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
|
||||
}
|
||||
|
||||
const float k_inv = is_fp8 ? 1.0f / static_cast<float>(k_scale) : 0.0f;
|
||||
const float v_inv = is_fp8 ? 1.0f / static_cast<float>(v_scale) : 0.0f;
|
||||
|
||||
const int64_t token_num = key.size(0);
|
||||
const int64_t key_token_num_stride = key.stride(0);
|
||||
const int64_t value_token_num_stride = value.stride(0);
|
||||
const int64_t head_num = value.size(1);
|
||||
const int64_t key_head_num_stride = key.stride(1);
|
||||
const int64_t value_head_num_stride = value.stride(1);
|
||||
const int64_t head_num = key.size(1);
|
||||
const int64_t head_dim = key.size(2);
|
||||
const int64_t num_blocks = key_cache.size(0);
|
||||
const int64_t num_blocks_stride = key_cache.stride(0);
|
||||
const int64_t cache_head_num_stride = key_cache.stride(1);
|
||||
const int64_t block_size = key_cache.size(2);
|
||||
const int64_t block_size_stride = key_cache.stride(2);
|
||||
const int64_t head_dim = key.size(-1);
|
||||
|
||||
cpu_attention::ISA isa_tag = [&]() {
|
||||
if (isa == "amx") {
|
||||
@@ -109,16 +134,24 @@ void cpu_attn_reshape_and_cache(
|
||||
}
|
||||
}();
|
||||
|
||||
if (is_fp8) {
|
||||
TORCH_CHECK(isa_tag == cpu_attention::ISA::AMX ||
|
||||
isa_tag == cpu_attention::ISA::VEC,
|
||||
"FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
|
||||
}
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
key.scalar_type(), "cpu_attn_reshape_and_cache", [&]() {
|
||||
CPU_ATTN_DISPATCH(head_dim, isa_tag, [&]() {
|
||||
CPU_ATTN_DISPATCH(head_dim, isa_tag, kv_cache_idx, [&]() {
|
||||
using kv_t = typename attn_impl::kv_cache_t;
|
||||
attn_impl::reshape_and_cache(
|
||||
key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
|
||||
key_cache.data_ptr<scalar_t>(), value_cache.data_ptr<scalar_t>(),
|
||||
slot_mapping.data_ptr<int64_t>(), token_num, key_token_num_stride,
|
||||
value_token_num_stride, head_num, key_head_num_stride,
|
||||
value_head_num_stride, num_blocks, num_blocks_stride,
|
||||
cache_head_num_stride, block_size, block_size_stride);
|
||||
reinterpret_cast<kv_t*>(key_cache.data_ptr()),
|
||||
reinterpret_cast<kv_t*>(value_cache.data_ptr()),
|
||||
slot_mapping.data_ptr<int64_t>(), token_num, key.stride(0),
|
||||
value.stride(0), head_num, key.stride(1), value.stride(1),
|
||||
num_blocks, num_blocks_stride, cache_head_num_stride, block_size,
|
||||
block_size_stride, k_inv, v_inv);
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -137,13 +170,26 @@ void cpu_attention_with_kv_cache(
|
||||
const int64_t sliding_window_left, const int64_t sliding_window_right,
|
||||
const torch::Tensor& block_table, // [num_tokens, max_block_num]
|
||||
const double softcap, const torch::Tensor& scheduler_metadata,
|
||||
const std::optional<torch::Tensor>& s_aux // [num_heads]
|
||||
) {
|
||||
const std::optional<torch::Tensor>& s_aux, // [num_heads]
|
||||
const double k_scale = 1.0, const double v_scale = 1.0,
|
||||
const std::string& kv_cache_dtype = "auto") {
|
||||
TORCH_CHECK_EQ(query.dim(), 3);
|
||||
TORCH_CHECK_EQ(query.stride(2), 1);
|
||||
TORCH_CHECK_EQ(key_cache.dim(), 4);
|
||||
TORCH_CHECK_EQ(value_cache.dim(), 4);
|
||||
|
||||
const int64_t kv_cache_idx =
|
||||
static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
|
||||
const bool is_fp8 = (kv_cache_idx != 0);
|
||||
if (is_fp8) {
|
||||
TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
|
||||
"key_cache must be uint8 for FP8 path");
|
||||
TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
|
||||
"value_cache must be uint8 for FP8 path");
|
||||
TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
|
||||
TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
|
||||
}
|
||||
|
||||
cpu_attention::AttentionInput input;
|
||||
input.metadata = reinterpret_cast<cpu_attention::AttentionMetadata*>(
|
||||
scheduler_metadata.data_ptr());
|
||||
@@ -165,25 +211,32 @@ void cpu_attention_with_kv_cache(
|
||||
input.block_table = block_table.data_ptr<int32_t>();
|
||||
input.alibi_slopes =
|
||||
alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
|
||||
// For now sink must be bf16
|
||||
input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
|
||||
input.scale = scale;
|
||||
input.causal = causal;
|
||||
input.sliding_window_left = sliding_window_left;
|
||||
input.sliding_window_right = sliding_window_right;
|
||||
if (input.causal) {
|
||||
// to make boundary calculation easier
|
||||
input.sliding_window_right = 0;
|
||||
}
|
||||
float softcap_fp32 = softcap;
|
||||
input.softcap = softcap_fp32;
|
||||
input.softcap = static_cast<float>(softcap);
|
||||
|
||||
if (is_fp8) {
|
||||
input.k_scale_fp8 = static_cast<float>(k_scale);
|
||||
input.v_scale_fp8 = static_cast<float>(v_scale);
|
||||
TORCH_CHECK(input.metadata->isa == cpu_attention::ISA::AMX ||
|
||||
input.metadata->isa == cpu_attention::ISA::VEC,
|
||||
"FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
|
||||
}
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
query.scalar_type(), "cpu_attention_with_kv_cache", [&]() {
|
||||
CPU_ATTN_DISPATCH(query.size(2), input.metadata->isa, [&]() {
|
||||
TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0);
|
||||
cpu_attention::AttentionMainLoop<attn_impl> mainloop;
|
||||
mainloop(&input);
|
||||
});
|
||||
CPU_ATTN_DISPATCH(
|
||||
query.size(2), input.metadata->isa, kv_cache_idx, [&]() {
|
||||
TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment,
|
||||
0);
|
||||
cpu_attention::AttentionMainLoop<attn_impl> mainloop;
|
||||
mainloop(&input);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
+171
-46
@@ -1,6 +1,7 @@
|
||||
#ifndef CPU_ATTN_AMX_HPP
|
||||
#define CPU_ATTN_AMX_HPP
|
||||
|
||||
#include "cpu_attn_fp8.hpp"
|
||||
#include "cpu_attn_impl.hpp"
|
||||
|
||||
namespace cpu_attention {
|
||||
@@ -21,9 +22,10 @@ typedef struct __tile_config {
|
||||
// 2-2-4 pattern, for 16 < m <= 32
|
||||
// TILE 0, 1: load A matrix, row num should be 16, m - 16
|
||||
// TILE 2, 3: load B matrix, row num should be 16
|
||||
// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, m - 16, m
|
||||
// - 16
|
||||
template <typename kv_cache_t>
|
||||
// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16,
|
||||
// m - 16, m - 16
|
||||
// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
|
||||
template <typename q_buffer_t, typename kv_cache_t>
|
||||
class TileGemm224 {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
@@ -42,13 +44,56 @@ class TileGemm224 {
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
class TileGemm224<c10::BFloat16> {
|
||||
// Dequantize one FP8 tile (AMX_TILE_ROW_NUM rows x 32 cols) to BF16.
|
||||
template <typename kv_cache_t>
|
||||
FORCE_INLINE void deq_tile_amx(const uint8_t* src, c10::BFloat16* dst) {
|
||||
for (int r = 0; r < AMX_TILE_ROW_NUM; ++r) {
|
||||
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e4m3_tag{})
|
||||
.save(dst + r * 32);
|
||||
} else {
|
||||
vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e5m2_tag{})
|
||||
.save(dst + r * 32);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// For FP8: dequant src into scratch and return scratch.
|
||||
// For BF16: return src directly (scratch is unused; the compiler elides it).
|
||||
template <typename kv_cache_t>
|
||||
FORCE_INLINE const c10::BFloat16* prepare_b_tile(const kv_cache_t* src,
|
||||
c10::BFloat16* scratch) {
|
||||
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
|
||||
deq_tile_amx<kv_cache_t>(reinterpret_cast<const uint8_t*>(src), scratch);
|
||||
return scratch;
|
||||
} else {
|
||||
return reinterpret_cast<const c10::BFloat16*>(src);
|
||||
}
|
||||
}
|
||||
|
||||
// Handles both BF16 and FP8 KV cache (2-2-4 pattern).
|
||||
template <typename kv_cache_t>
|
||||
class TileGemm224<c10::BFloat16, kv_cache_t> {
|
||||
static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
|
||||
"kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
|
||||
|
||||
static constexpr bool fp8_kv =
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
|
||||
|
||||
static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
// BF16 path: scratch_elems=1 so the scratch array is eliminated by the
|
||||
// compiler.
|
||||
static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
|
||||
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
c10::BFloat16* __restrict__ a_tile,
|
||||
c10::BFloat16* __restrict__ b_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,
|
||||
@@ -56,6 +101,7 @@ class TileGemm224<c10::BFloat16> {
|
||||
const bool accum_c) {
|
||||
const int32_t k_times =
|
||||
dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
|
||||
|
||||
c10::BFloat16* __restrict__ a_tile_0 = a_tile;
|
||||
c10::BFloat16* __restrict__ a_tile_1 = a_tile + lda * AMX_TILE_ROW_NUM;
|
||||
const int64_t a_tile_stride = [&]() {
|
||||
@@ -70,8 +116,8 @@ class TileGemm224<c10::BFloat16> {
|
||||
}
|
||||
}();
|
||||
|
||||
c10::BFloat16* __restrict__ b_tile_2 = b_tile;
|
||||
c10::BFloat16* __restrict__ b_tile_3 = [&]() {
|
||||
kv_cache_t* __restrict__ b_tile_2 = b_tile;
|
||||
kv_cache_t* __restrict__ b_tile_3 = [&]() {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
// k_cache is prepacked
|
||||
return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
|
||||
@@ -106,11 +152,16 @@ class TileGemm224<c10::BFloat16> {
|
||||
_tile_zero(7);
|
||||
}
|
||||
|
||||
alignas(64) c10::BFloat16 scratch_2[scratch_elems];
|
||||
alignas(64) c10::BFloat16 scratch_3[scratch_elems];
|
||||
for (int32_t k = 0; k < k_times; ++k) {
|
||||
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
|
||||
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
|
||||
|
||||
_tile_loadd(0, a_tile_0, a_tile_stride);
|
||||
_tile_stream_loadd(2, b_tile_2, b_tile_stride);
|
||||
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_tile_stride);
|
||||
_tile_dpbf16ps(4, 0, 2);
|
||||
_tile_stream_loadd(3, b_tile_3, b_tile_stride);
|
||||
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_tile_stride);
|
||||
_tile_dpbf16ps(5, 0, 3);
|
||||
_tile_loadd(1, a_tile_1, a_tile_stride);
|
||||
_tile_dpbf16ps(6, 1, 2);
|
||||
@@ -154,13 +205,13 @@ class TileGemm224<c10::BFloat16> {
|
||||
};
|
||||
|
||||
// 1-2-2 pattern, for 0 < m <= 16
|
||||
// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should be
|
||||
// m, m
|
||||
// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row
|
||||
// num should be 16
|
||||
// TILE 6, 7, (6, 7): store results C matrix, row num should be
|
||||
// m
|
||||
template <typename kv_cache_t>
|
||||
// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should
|
||||
// be m, m
|
||||
// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row num
|
||||
// should be 16
|
||||
// TILE 6, 7: store results C matrix, row num should be m
|
||||
// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
|
||||
template <typename q_buffer_t, typename kv_cache_t>
|
||||
class TileGemm122 {
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
@@ -179,13 +230,26 @@ class TileGemm122 {
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
class TileGemm122<c10::BFloat16> {
|
||||
// Handles both BF16 and FP8 KV cache (1-2-2 pattern).
|
||||
template <typename kv_cache_t>
|
||||
class TileGemm122<c10::BFloat16, kv_cache_t> {
|
||||
static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
|
||||
"kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
|
||||
|
||||
static constexpr bool fp8_kv =
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
|
||||
|
||||
static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
|
||||
|
||||
public:
|
||||
template <AttentionGemmPhase phase, int32_t k_size>
|
||||
FORCE_INLINE static void gemm(const int32_t m_size,
|
||||
c10::BFloat16* __restrict__ a_tile,
|
||||
c10::BFloat16* __restrict__ b_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,
|
||||
@@ -215,21 +279,19 @@ class TileGemm122<c10::BFloat16> {
|
||||
}
|
||||
}();
|
||||
|
||||
c10::BFloat16* __restrict__ b_tile_2 = b_tile;
|
||||
c10::BFloat16* __restrict__ b_tile_3 = [&]() {
|
||||
kv_cache_t* __restrict__ b_tile_2 = b_tile;
|
||||
kv_cache_t* __restrict__ b_tile_3 = [&]() {
|
||||
if constexpr (phase == AttentionGemmPhase::QK) {
|
||||
// k_cache is prepacked
|
||||
return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
|
||||
} else if constexpr (phase == AttentionGemmPhase::PV) {
|
||||
// v_cache is prepacked
|
||||
return b_tile + (block_size * AMX_TILE_ROW_BYTES / 4);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unreachable");
|
||||
}
|
||||
}();
|
||||
c10::BFloat16* __restrict__ b_tile_4 =
|
||||
kv_cache_t* __restrict__ b_tile_4 =
|
||||
b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
c10::BFloat16* __restrict__ b_tile_5 =
|
||||
kv_cache_t* __restrict__ b_tile_5 =
|
||||
b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
|
||||
int64_t b_stride = AMX_TILE_ROW_BYTES;
|
||||
|
||||
@@ -250,16 +312,25 @@ class TileGemm122<c10::BFloat16> {
|
||||
_tile_zero(7);
|
||||
}
|
||||
|
||||
alignas(64) c10::BFloat16 scratch_2[scratch_elems];
|
||||
alignas(64) c10::BFloat16 scratch_3[scratch_elems];
|
||||
alignas(64) c10::BFloat16 scratch_4[scratch_elems];
|
||||
alignas(64) c10::BFloat16 scratch_5[scratch_elems];
|
||||
for (int32_t k = 0; k < k_group_times; ++k) {
|
||||
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
|
||||
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
|
||||
const c10::BFloat16* load_4 = prepare_b_tile(b_tile_4, scratch_4);
|
||||
const c10::BFloat16* load_5 = prepare_b_tile(b_tile_5, scratch_5);
|
||||
|
||||
_tile_loadd(0, a_tile_0, a_tile_stride);
|
||||
_tile_stream_loadd(2, b_tile_2, b_stride);
|
||||
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
|
||||
_tile_dpbf16ps(6, 0, 2);
|
||||
_tile_stream_loadd(3, b_tile_3, b_stride);
|
||||
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
|
||||
_tile_dpbf16ps(7, 0, 3);
|
||||
_tile_loadd(1, a_tile_1, a_tile_stride);
|
||||
_tile_stream_loadd(4, b_tile_4, b_stride);
|
||||
_tile_stream_loadd(4, const_cast<c10::BFloat16*>(load_4), b_stride);
|
||||
_tile_dpbf16ps(6, 1, 4);
|
||||
_tile_stream_loadd(5, b_tile_5, b_stride);
|
||||
_tile_stream_loadd(5, const_cast<c10::BFloat16*>(load_5), b_stride);
|
||||
_tile_dpbf16ps(7, 1, 5);
|
||||
|
||||
// update ptrs
|
||||
@@ -279,10 +350,13 @@ class TileGemm122<c10::BFloat16> {
|
||||
}
|
||||
|
||||
if (has_tail) {
|
||||
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
|
||||
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
|
||||
|
||||
_tile_loadd(0, a_tile_0, a_tile_stride);
|
||||
_tile_stream_loadd(2, b_tile_2, b_stride);
|
||||
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
|
||||
_tile_dpbf16ps(6, 0, 2);
|
||||
_tile_stream_loadd(3, b_tile_3, b_stride);
|
||||
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
|
||||
_tile_dpbf16ps(7, 0, 3);
|
||||
}
|
||||
|
||||
@@ -302,21 +376,25 @@ class TileGemm122<c10::BFloat16> {
|
||||
_tile_loadconfig(&config);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t, int64_t head_dim>
|
||||
class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
static constexpr bool fp8_kv =
|
||||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
|
||||
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = scalar_t;
|
||||
using kv_cache_t = scalar_t;
|
||||
using kv_cache_t = kv_cache_scalar_t;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = scalar_t;
|
||||
|
||||
constexpr static int64_t BlockSizeAlignment =
|
||||
AMX_TILE_ROW_BYTES /
|
||||
sizeof(kv_cache_t); // KV token num unit of QK and PV phases
|
||||
32; // AMX_TILE_ROW_NUM = 16 tokens/tile; 32 = 2 tiles
|
||||
constexpr static int64_t HeadDimAlignment =
|
||||
2 * (AMX_TILE_ROW_BYTES / 4); // headdim num unit of PV phase
|
||||
constexpr static int64_t MaxQHeadNumPerIteration = 32;
|
||||
@@ -324,6 +402,9 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
|
||||
constexpr static ISA ISAType = ISA::AMX;
|
||||
constexpr static bool scale_on_logits = true;
|
||||
|
||||
float k_scale = 1.0f;
|
||||
float v_scale = 1.0f;
|
||||
|
||||
public:
|
||||
AttentionImpl() : current_q_head_num_(0) {
|
||||
// Use all columns in AMX tiles
|
||||
@@ -332,21 +413,50 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
|
||||
|
||||
~AttentionImpl() { _tile_release(); }
|
||||
|
||||
void init_from_input(const AttentionInput* input) {
|
||||
if constexpr (fp8_kv) {
|
||||
k_scale = input->k_scale_fp8;
|
||||
v_scale = input->v_scale_fp8;
|
||||
}
|
||||
}
|
||||
|
||||
float get_output_v_scale() const noexcept {
|
||||
if constexpr (fp8_kv) {
|
||||
// AMX dequant places FP8 payload into a BF16 field (exponent bias 127).
|
||||
// Correction = 2^(127 - FP8_bias): E4M3 bias=7 → 2^120, E5M2 bias=15 →
|
||||
// 2^112.
|
||||
constexpr float bias =
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
|
||||
return v_scale * bias;
|
||||
}
|
||||
return 1.0f;
|
||||
}
|
||||
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
if constexpr (fp8_kv) {
|
||||
// Same bias correction as get_output_v_scale: AMX FP8→BF16 dequant
|
||||
// shifts the exponent bias from FP8 to BF16 (127), so we multiply by
|
||||
// 2^(127-FP8_bias) to recover the true value. E4M3: 2^120, E5M2: 2^112.
|
||||
const float bias =
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
|
||||
scale *= k_scale * bias;
|
||||
}
|
||||
if (q_head_num > AMX_TILE_ROW_NUM) {
|
||||
if (q_head_num != current_q_head_num_) {
|
||||
current_q_head_num_ = q_head_num;
|
||||
TileGemm224<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
|
||||
TileGemm224<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
|
||||
amx_tile_config_);
|
||||
}
|
||||
attention<TileGemm224<kv_cache_t>> attention_iteration;
|
||||
attention<TileGemm224<q_buffer_t, kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
} else {
|
||||
if (q_head_num != current_q_head_num_) {
|
||||
current_q_head_num_ = q_head_num;
|
||||
TileGemm122<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
|
||||
TileGemm122<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
|
||||
amx_tile_config_);
|
||||
}
|
||||
attention<TileGemm122<kv_cache_t>> attention_iteration;
|
||||
attention<TileGemm122<q_buffer_t, kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
}
|
||||
@@ -411,13 +521,26 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
|
||||
// reshape KV to AMX friendly layout
|
||||
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,
|
||||
kv_cache_t* __restrict__ key_cache, kv_cache_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 int64_t block_size, const int64_t block_size_stride,
|
||||
const float k_inv = 0.0f, const float v_inv = 0.0f) {
|
||||
if constexpr (fp8_kv) {
|
||||
constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
|
||||
reshape_and_cache_fp8_amx_impl<scalar_t, qfn>(
|
||||
key, value, reinterpret_cast<uint8_t*>(key_cache),
|
||||
reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
|
||||
head_num, head_dim, block_size, key_token_num_stride,
|
||||
key_head_num_stride, value_token_num_stride, value_head_num_stride,
|
||||
num_blocks_stride, cache_head_num_stride, num_blocks_stride,
|
||||
cache_head_num_stride, k_inv, v_inv);
|
||||
return;
|
||||
}
|
||||
|
||||
// For AMX 2D tiles, size of each line is 64 bytes
|
||||
constexpr int64_t amx_tile_row_size = AMX_TILE_ROW_BYTES;
|
||||
// For AMX B matrix, N always is 16
|
||||
@@ -426,6 +549,9 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
|
||||
// For now suppose block_size is divisible by amx_tile_column_num
|
||||
TORCH_CHECK_EQ(block_size % amx_b_tile_k_size, 0);
|
||||
|
||||
scalar_t* __restrict__ kc = reinterpret_cast<scalar_t*>(key_cache);
|
||||
scalar_t* __restrict__ vc = reinterpret_cast<scalar_t*>(value_cache);
|
||||
|
||||
#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) {
|
||||
@@ -453,8 +579,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
|
||||
constexpr int64_t quadword_num_per_group =
|
||||
token_num_per_group * quadword_num;
|
||||
int32_t* key_cache_start_ptr =
|
||||
reinterpret_cast<int32_t*>(key_cache +
|
||||
block_idx * num_blocks_stride +
|
||||
reinterpret_cast<int32_t*>(kc + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride) +
|
||||
group_idx * quadword_num_per_group + group_offset;
|
||||
|
||||
@@ -483,7 +608,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
|
||||
token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
scalar_t* value_cache_start_ptr =
|
||||
value_cache + block_idx * num_blocks_stride +
|
||||
vc + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride +
|
||||
sub_group_idx * token_num_per_sub_group * amx_b_tile_n_size +
|
||||
sub_group_offset;
|
||||
|
||||
@@ -0,0 +1,214 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#pragma once
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <limits>
|
||||
#include <type_traits>
|
||||
|
||||
#include "cpu/utils.hpp"
|
||||
|
||||
typedef uint32_t __attribute__((__may_alias__)) u32_alias_t;
|
||||
typedef uint16_t __attribute__((__may_alias__)) u16_alias_t;
|
||||
typedef float __attribute__((__may_alias__)) f32_alias_t;
|
||||
|
||||
// Reference scalar dequant — used to verify vectorized AMX dequant.
|
||||
inline float fp8e4m3_to_float_scalar(uint8_t b, float scale) noexcept {
|
||||
// NaN encoding in E4M3
|
||||
if ((b & 0x7F) == 0x7F) return std::numeric_limits<float>::quiet_NaN();
|
||||
uint32_t b_u32 = static_cast<uint32_t>(b);
|
||||
uint32_t sign = (b_u32 & 0x80) << 24;
|
||||
uint32_t payload = (b_u32 & 0x7F) << 20;
|
||||
uint32_t bits = sign | payload;
|
||||
float b_f32_unscaled = *reinterpret_cast<const f32_alias_t*>(&bits);
|
||||
float b_f32_scaled = b_f32_unscaled * scale * 0x1p120f;
|
||||
return b_f32_scaled;
|
||||
}
|
||||
|
||||
inline uint8_t float_to_fp8e4m3_scalar(float v, float inv_scale) noexcept {
|
||||
v *= inv_scale;
|
||||
constexpr float fp8_max = 448.0f;
|
||||
v = std::max(-fp8_max, std::min(fp8_max, v));
|
||||
if (v == 0.0f) return 0;
|
||||
|
||||
// Inverse mapping of fp8e4m3_to_float_scalar: shift the effective exponent
|
||||
// bias from fp32 (127) back to fp8 e4m3 (7), then pack sign|payload.
|
||||
float v_f32_unscaled = v * 0x1p-120f;
|
||||
uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v_f32_unscaled);
|
||||
uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
|
||||
uint8_t payload = static_cast<uint8_t>((bits >> 20) & 0x7F);
|
||||
if (payload == 0) return sign;
|
||||
payload = std::min<uint8_t>(payload, 0x7E); // keep 0x7F as NaN encoding
|
||||
return static_cast<uint8_t>(sign | payload);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// AMX reshape impl — parameterised on the quantisation function.
|
||||
// Writes key/value into uint8 FP8 KV cache using the AMX tile-friendly layout.
|
||||
// K: halfword-packed (2 FP8 per uint16, token_num_per_group=16).
|
||||
// V: sub-group packing (token_num_per_sub_group=2, head_elems_per_group=16).
|
||||
// block_size must be divisible by 32.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
|
||||
inline void reshape_and_cache_fp8_amx_impl(
|
||||
const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
|
||||
uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
|
||||
int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
|
||||
int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
|
||||
int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
|
||||
float v_inv) {
|
||||
constexpr int64_t token_num_per_group = 16; // AMX_TILE_ROW_NUM
|
||||
const int64_t halfword_num = head_dim / 2; // 2 FP8 per uint16
|
||||
const int64_t halfword_num_per_group = token_num_per_group * halfword_num;
|
||||
constexpr int64_t head_elems_per_group = 16;
|
||||
constexpr int64_t token_num_per_sub_group = 2; // = 4 / sizeof(BF16)
|
||||
const int64_t group_num = head_dim / head_elems_per_group;
|
||||
const int64_t group_size = block_size * head_elems_per_group;
|
||||
|
||||
#pragma omp parallel for collapse(2) schedule(static)
|
||||
for (int64_t tok = 0; tok < token_num; ++tok) {
|
||||
for (int64_t h = 0; h < head_num; ++h) {
|
||||
const int64_t slot = slot_ptr[tok];
|
||||
if (slot < 0) continue;
|
||||
const int64_t block_idx = slot / block_size;
|
||||
const int64_t block_offset = slot % block_size;
|
||||
|
||||
// Key: halfword-packed, 2 FP8 per uint16
|
||||
{
|
||||
const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
|
||||
const int64_t group_idx = block_offset / token_num_per_group;
|
||||
const int64_t group_offset = block_offset % token_num_per_group;
|
||||
uint16_t* kdst =
|
||||
reinterpret_cast<uint16_t*>(key_cache_ptr + block_idx * kc_stride0 +
|
||||
h * kc_stride1) +
|
||||
group_idx * halfword_num_per_group + group_offset;
|
||||
for (int64_t j = 0; j < halfword_num; ++j) {
|
||||
uint8_t fp8_0 = quant_fn(static_cast<float>(ksrc[j * 2]), k_inv);
|
||||
uint8_t fp8_1 = quant_fn(static_cast<float>(ksrc[j * 2 + 1]), k_inv);
|
||||
uint8_t bytes[2] = {fp8_0, fp8_1};
|
||||
uint16_t hw = *reinterpret_cast<const u16_alias_t*>(bytes);
|
||||
kdst[j * token_num_per_group] = hw;
|
||||
}
|
||||
}
|
||||
|
||||
// Value: sub-group packing (token_num_per_sub_group = 2)
|
||||
{
|
||||
const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
|
||||
const int64_t sub_group_idx = block_offset / token_num_per_sub_group;
|
||||
const int64_t sub_group_offset = block_offset % token_num_per_sub_group;
|
||||
uint8_t* vdst =
|
||||
value_cache_ptr + block_idx * vc_stride0 + h * vc_stride1 +
|
||||
sub_group_idx * token_num_per_sub_group * head_elems_per_group +
|
||||
sub_group_offset;
|
||||
for (int64_t i = 0; i < group_num; ++i) {
|
||||
for (int64_t j = 0; j < head_elems_per_group; ++j)
|
||||
vdst[j * token_num_per_sub_group] =
|
||||
quant_fn(static_cast<float>(vsrc[j]), v_inv);
|
||||
vsrc += head_elems_per_group;
|
||||
vdst += group_size;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// FP8 E5M2 scalar helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Reference scalar dequant — used to verify vectorized AMX dequant.
|
||||
// FP8 E5M2: s[7] e[6:2] m[1:0], exponent bias = 15 (same as FP16).
|
||||
// Byte b → FP16 bits = b << 8 (no bias correction needed).
|
||||
inline float fp8e5m2_to_float_scalar(uint8_t b, float scale) noexcept {
|
||||
const uint8_t exp_bits = (b >> 2) & 0x1F;
|
||||
const uint8_t mant_bits = b & 0x03;
|
||||
// NaN: exp=11111, mant!=00
|
||||
if (exp_bits == 0x1F && mant_bits != 0)
|
||||
return std::numeric_limits<float>::quiet_NaN();
|
||||
const uint32_t sign = static_cast<uint32_t>(b & 0x80) << 24;
|
||||
if (exp_bits == 0x1F)
|
||||
return sign ? -std::numeric_limits<float>::infinity()
|
||||
: std::numeric_limits<float>::infinity();
|
||||
if (exp_bits == 0) { // subnormal: (-1)^s * 2^-14 * mant/4
|
||||
if (mant_bits == 0) return 0.0f;
|
||||
float v = mant_bits * 0x1p-16f;
|
||||
return (sign ? -v : v) * scale;
|
||||
}
|
||||
// Normal: FP32 exp = exp5 - 15 + 127, mantissa top 2 bits
|
||||
uint32_t fp32_bits = sign |
|
||||
((static_cast<uint32_t>(exp_bits) - 15 + 127) << 23) |
|
||||
(static_cast<uint32_t>(mant_bits) << 21);
|
||||
float val = *reinterpret_cast<const f32_alias_t*>(&fp32_bits);
|
||||
return val * scale;
|
||||
}
|
||||
|
||||
inline uint8_t float_to_fp8e5m2_scalar(float v, float inv_scale) noexcept {
|
||||
v *= inv_scale;
|
||||
constexpr float fp8_e5m2_max = 57344.0f;
|
||||
v = std::max(-fp8_e5m2_max, std::min(fp8_e5m2_max, v));
|
||||
if (v == 0.0f) return 0;
|
||||
uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v);
|
||||
const uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
|
||||
const int32_t exp_fp32 = static_cast<int32_t>((bits >> 23) & 0xFF) - 127;
|
||||
const uint8_t mant2 = static_cast<uint8_t>((bits >> 21) & 0x03);
|
||||
if (exp_fp32 < -14) { // subnormal in E5M2
|
||||
const int shift = -14 - exp_fp32;
|
||||
if (shift + 21 >= 32)
|
||||
return sign; // underflow: too small for E5M2 subnormal
|
||||
const uint32_t m = (0x800000u | (bits & 0x7FFFFFu)) >> (shift + 21);
|
||||
return sign | static_cast<uint8_t>(std::min<uint32_t>(m, 3u));
|
||||
}
|
||||
const uint8_t exp5 = static_cast<uint8_t>(exp_fp32 + 15);
|
||||
return sign | (exp5 << 2) | mant2;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Select the FP8 quant function at compile time based on kv_cache_t.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename kv_cache_t>
|
||||
constexpr auto select_fp8_quant_fn() {
|
||||
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>)
|
||||
return float_to_fp8e5m2_scalar;
|
||||
else
|
||||
return float_to_fp8e4m3_scalar;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// VEC reshape impl — parameterised on the quantisation function.
|
||||
// Writes key (column-major) and value (row-major) into uint8 FP8 KV cache.
|
||||
// The pragma omp must live outside VLLM_DISPATCH_FLOATING_TYPES because
|
||||
// #pragma cannot appear inside variadic macro arguments.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
|
||||
inline void reshape_and_cache_fp8_vec_impl(
|
||||
const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
|
||||
uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
|
||||
int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
|
||||
int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
|
||||
int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
|
||||
float v_inv) {
|
||||
#pragma omp parallel for collapse(2) schedule(static)
|
||||
for (int64_t tok = 0; tok < token_num; ++tok) {
|
||||
for (int64_t h = 0; h < head_num; ++h) {
|
||||
const int64_t slot = slot_ptr[tok];
|
||||
if (slot < 0) continue;
|
||||
const int64_t block_idx = slot / block_size;
|
||||
const int64_t block_offset = slot % block_size;
|
||||
|
||||
// Key layout: column-major within block
|
||||
const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
|
||||
uint8_t* kdst = key_cache_ptr + block_idx * kc_stride0 + h * kc_stride1 +
|
||||
block_offset;
|
||||
for (int64_t i = 0; i < head_dim; ++i)
|
||||
kdst[i * block_size] = quant_fn(static_cast<float>(ksrc[i]), k_inv);
|
||||
|
||||
// Value layout: row-major within block (contiguous head_dim bytes)
|
||||
const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
|
||||
uint8_t* vdst = value_cache_ptr + block_idx * vc_stride0 +
|
||||
h * vc_stride1 + block_offset * head_dim;
|
||||
for (int64_t i = 0; i < head_dim; ++i)
|
||||
vdst[i] = quant_fn(static_cast<float>(vsrc[i]), v_inv);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -14,8 +14,22 @@
|
||||
namespace cpu_attention {
|
||||
enum class ISA { AMX, VEC, VEC16, NEON, VXE };
|
||||
|
||||
template <ISA isa, typename scalar_t, int64_t head_dim>
|
||||
class AttentionImpl {};
|
||||
// Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
|
||||
enum class Fp8KVCacheDataType {
|
||||
kAuto = 0,
|
||||
kFp8E4M3 = 1,
|
||||
kFp8E5M2 = 2,
|
||||
};
|
||||
|
||||
struct AttentionInput;
|
||||
|
||||
template <ISA isa, typename scalar_t, int64_t head_dim,
|
||||
typename kv_cache_scalar_t = scalar_t>
|
||||
class AttentionImpl {
|
||||
public:
|
||||
void init_from_input(const AttentionInput*) {}
|
||||
float get_output_v_scale() const noexcept { return 1.0f; }
|
||||
};
|
||||
|
||||
struct AttentionWorkItemGroup {
|
||||
int32_t req_id;
|
||||
@@ -780,6 +794,9 @@ struct AttentionInput {
|
||||
int32_t sliding_window_left;
|
||||
int32_t sliding_window_right;
|
||||
float softcap;
|
||||
// FP8 KV cache scales (used by FP8 attention implementations)
|
||||
float k_scale_fp8 = 1.0f;
|
||||
float v_scale_fp8 = 1.0f;
|
||||
};
|
||||
|
||||
#define DEFINE_CPU_ATTENTION_PARAMS \
|
||||
@@ -1374,6 +1391,13 @@ class AttentionMainLoop {
|
||||
}
|
||||
|
||||
attention_impl_t attn_impl;
|
||||
constexpr bool fp8_kv = std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
|
||||
float output_v_scale = 1.0f;
|
||||
if constexpr (fp8_kv) {
|
||||
attn_impl.init_from_input(input);
|
||||
output_v_scale = attn_impl.get_output_v_scale();
|
||||
}
|
||||
|
||||
// general information
|
||||
const int32_t q_head_num = input->num_heads;
|
||||
@@ -1753,7 +1777,7 @@ class AttentionMainLoop {
|
||||
reinterpret_cast<query_t*>(input->output) +
|
||||
output_buffer_offset,
|
||||
sum_buffer, actual_q_heads_per_kv,
|
||||
actual_q_token_num, q_head_num);
|
||||
actual_q_token_num, q_head_num, output_v_scale);
|
||||
} else {
|
||||
const int32_t stride =
|
||||
actual_q_heads_per_kv * split_kv_q_token_num_threshold;
|
||||
@@ -1823,7 +1847,7 @@ class AttentionMainLoop {
|
||||
split_output_buffer,
|
||||
reinterpret_cast<query_t*>(input->output) + output_buffer_offset,
|
||||
split_sum_buffer, actual_q_heads_per_kv, curr_output_token_num,
|
||||
q_head_num);
|
||||
q_head_num, output_v_scale);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1947,8 +1971,8 @@ class AttentionMainLoop {
|
||||
query_t* __restrict__ curr_output_buffer,
|
||||
float* __restrict__ sum_buffer,
|
||||
const int32_t q_heads_per_kv,
|
||||
const int32_t actual_q_token_num,
|
||||
const int32_t q_head_num) {
|
||||
const int32_t actual_q_token_num, const int32_t q_head_num,
|
||||
const float v_scale = 1.0f) {
|
||||
// final output
|
||||
using output_vec_t = typename VecTypeTrait<query_t>::vec_t;
|
||||
|
||||
@@ -1962,7 +1986,7 @@ class AttentionMainLoop {
|
||||
curr_partial_output_buffer;
|
||||
query_t* __restrict__ curr_output_buffer_iter = curr_output_buffer;
|
||||
for (int32_t head_idx = 0; head_idx < q_heads_per_kv; ++head_idx) {
|
||||
vec_op::FP32Vec16 inv_sum_scale_vec(1.0 / *curr_sum_buffer);
|
||||
vec_op::FP32Vec16 inv_sum_scale_vec(v_scale / *curr_sum_buffer);
|
||||
|
||||
for (int32_t i = 0; i < group_num_per_head; ++i) {
|
||||
vec_op::FP32Vec16 vec(curr_partial_output_buffer_iter);
|
||||
|
||||
@@ -248,8 +248,8 @@ class TileGemmNeonFMLA {
|
||||
} // namespace
|
||||
|
||||
// this is similar to "ISA::VEC" at the moment
|
||||
template <typename scalar_t, int64_t head_dim>
|
||||
class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
@@ -343,7 +343,8 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
|
||||
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 int64_t block_size, const int64_t block_size_stride,
|
||||
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
|
||||
#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) {
|
||||
@@ -388,7 +389,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
|
||||
#ifdef ARM_BF16_SUPPORT
|
||||
// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
|
||||
template <int64_t head_dim>
|
||||
class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim>
|
||||
class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim, c10::BFloat16>
|
||||
: public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
|
||||
head_dim> {};
|
||||
#endif
|
||||
|
||||
@@ -602,7 +602,8 @@ class AttentionImplNEONBFMMLA {
|
||||
[[maybe_unused]] const int64_t num_blocks,
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size,
|
||||
[[maybe_unused]] const int64_t block_size_stride) {
|
||||
[[maybe_unused]] const int64_t block_size_stride,
|
||||
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
|
||||
const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
|
||||
const int64_t v_pair_stride =
|
||||
(block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
|
||||
|
||||
+105
-28
@@ -1,11 +1,37 @@
|
||||
#ifndef CPU_ATTN_VEC_HPP
|
||||
#define CPU_ATTN_VEC_HPP
|
||||
|
||||
#include "cpu_attn_fp8.hpp"
|
||||
#include "cpu_attn_impl.hpp"
|
||||
|
||||
namespace cpu_attention {
|
||||
|
||||
namespace {
|
||||
|
||||
// Load 32 kv_cache_t elements starting at ptr and return them as two FP32Vec16s
|
||||
// covering the lower 16 and upper 16 positions.
|
||||
// For FP8: both halves come from a single BF16Vec32 dequant of 32 bytes.
|
||||
// For BF16/FP16/FP32: two separate vector loads at ptr and ptr+16.
|
||||
template <typename kv_cache_t>
|
||||
FORCE_INLINE std::pair<vec_op::FP32Vec16, vec_op::FP32Vec16> load_b_pair_vec(
|
||||
const kv_cache_t* ptr) {
|
||||
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
|
||||
// BF16 container, but values are in the FP16 exponent range (bias 15 not
|
||||
// 127).
|
||||
vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
|
||||
vec_op::fp8_e4m3_tag{});
|
||||
return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
|
||||
} else if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
|
||||
vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
|
||||
vec_op::fp8_e5m2_tag{});
|
||||
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))};
|
||||
}
|
||||
}
|
||||
|
||||
// 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
|
||||
template <typename kv_cache_t>
|
||||
class TileGemm82 {
|
||||
@@ -54,10 +80,7 @@ class TileGemm82 {
|
||||
const int32_t block_size, const int32_t dynamic_k_size,
|
||||
const bool accum_c) {
|
||||
static_assert(0 < M && M <= 8);
|
||||
using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
|
||||
|
||||
kv_cache_t* __restrict__ curr_b_0 = b_tile;
|
||||
kv_cache_t* __restrict__ curr_b_1 = b_tile + 16;
|
||||
float* __restrict__ curr_c_0 = c_tile;
|
||||
float* __restrict__ curr_c_1 = c_tile + 16;
|
||||
|
||||
@@ -76,16 +99,14 @@ class TileGemm82 {
|
||||
}
|
||||
|
||||
float* __restrict__ curr_a = a_tile;
|
||||
kv_cache_t* __restrict__ curr_b = b_tile;
|
||||
|
||||
for (int32_t k = 0; k < dynamic_k_size; ++k) {
|
||||
load_vec_t b_0_reg(curr_b_0);
|
||||
vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
|
||||
load_vec_t b_1_reg(curr_b_1);
|
||||
vec_op::FP32Vec16 fp32_b_1_reg(b_1_reg);
|
||||
auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
|
||||
|
||||
float* __restrict__ curr_m_a = curr_a;
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
float v = *curr_m_a;
|
||||
vec_op::FP32Vec16 a_reg(v);
|
||||
vec_op::FP32Vec16 a_reg(*curr_m_a);
|
||||
c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
|
||||
c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
|
||||
|
||||
@@ -95,8 +116,7 @@ class TileGemm82 {
|
||||
|
||||
// update
|
||||
curr_a += 1;
|
||||
curr_b_0 += ldb;
|
||||
curr_b_1 += ldb;
|
||||
curr_b += ldb;
|
||||
}
|
||||
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
@@ -109,15 +129,20 @@ class TileGemm82 {
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
// This is a general but naive implementation based on vector instructions
|
||||
template <typename scalar_t, int64_t head_dim>
|
||||
class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
static constexpr bool fp8_kv =
|
||||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
|
||||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
|
||||
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
using kv_cache_t = scalar_t;
|
||||
using kv_cache_t = kv_cache_scalar_t;
|
||||
using logits_buffer_t = float;
|
||||
using partial_output_buffer_t = float;
|
||||
using prob_buffer_t = float;
|
||||
@@ -129,11 +154,45 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
|
||||
constexpr static int64_t MaxQHeadNumPerIteration = 8;
|
||||
constexpr static int64_t HeadDim = head_dim;
|
||||
constexpr static ISA ISAType = ISA::VEC;
|
||||
constexpr static bool scale_on_logits = false; // apply scale on q_buffer
|
||||
constexpr static bool scale_on_logits = fp8_kv;
|
||||
|
||||
float k_scale = 1.0f;
|
||||
float v_scale = 1.0f;
|
||||
|
||||
public:
|
||||
void init_from_input(const AttentionInput* input) {
|
||||
if constexpr (fp8_kv) {
|
||||
k_scale = input->k_scale_fp8;
|
||||
v_scale = input->v_scale_fp8;
|
||||
}
|
||||
}
|
||||
|
||||
float get_output_v_scale() const noexcept {
|
||||
if constexpr (fp8_kv) {
|
||||
// VEC dequant unpacks FP8 into a pseudo-FP16 layout (exponent bias 15).
|
||||
// E4M3 (bias=7) needs correction 2^(15-7) = 2^8; E5M2 bias matches FP16
|
||||
// so no correction.
|
||||
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
|
||||
return v_scale;
|
||||
} else {
|
||||
return v_scale * 0x1p8f;
|
||||
}
|
||||
}
|
||||
return 1.0f;
|
||||
}
|
||||
|
||||
template <template <typename tile_gemm_t> typename attention>
|
||||
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
|
||||
if constexpr (fp8_kv) {
|
||||
// Same bias correction as get_output_v_scale: VEC FP8→pseudo-FP16 dequant
|
||||
// uses bias 15; E4M3 (bias=7) needs ×2^8, E5M2 (bias=15) needs no
|
||||
// correction.
|
||||
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
|
||||
scale *= k_scale;
|
||||
} else {
|
||||
scale *= k_scale * 0x1p8f;
|
||||
}
|
||||
}
|
||||
attention<TileGemm82<kv_cache_t>> attention_iteration;
|
||||
attention_iteration(CPU_ATTENTION_PARAMS);
|
||||
}
|
||||
@@ -161,17 +220,19 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
|
||||
// row-major
|
||||
}
|
||||
|
||||
// Copy q to q_buffer and cast it to fp32
|
||||
static void copy_q_heads_tile(
|
||||
scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size]
|
||||
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) {
|
||||
// Copy q to q_buffer and cast it to fp32.
|
||||
// FP8: QK scale is folded into execute_attention; copy Q unscaled here.
|
||||
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) {
|
||||
static_assert(head_dim % 16 == 0);
|
||||
constexpr int32_t unroll_size = head_dim / 16;
|
||||
using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
|
||||
|
||||
vec_op::FP32Vec16 scale_vec(scale);
|
||||
const float effective_scale = fp8_kv ? 1.0f : scale;
|
||||
vec_op::FP32Vec16 scale_vec(effective_scale);
|
||||
for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
|
||||
for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
|
||||
scalar_t* __restrict__ curr_q =
|
||||
@@ -196,13 +257,26 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
|
||||
// reshape K as column-major and V as row-major
|
||||
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,
|
||||
kv_cache_t* __restrict__ key_cache, kv_cache_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 int64_t block_size, const int64_t block_size_stride,
|
||||
const float k_inv = 0.0f, const float v_inv = 0.0f) {
|
||||
if constexpr (fp8_kv) {
|
||||
constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
|
||||
reshape_and_cache_fp8_vec_impl<scalar_t, qfn>(
|
||||
key, value, reinterpret_cast<uint8_t*>(key_cache),
|
||||
reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
|
||||
head_num, head_dim, block_size, key_token_num_stride,
|
||||
key_head_num_stride, value_token_num_stride, value_head_num_stride,
|
||||
num_blocks_stride, cache_head_num_stride, num_blocks_stride,
|
||||
cache_head_num_stride, k_inv, v_inv);
|
||||
return;
|
||||
}
|
||||
|
||||
#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) {
|
||||
@@ -220,8 +294,9 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
|
||||
token_idx * key_token_num_stride +
|
||||
head_idx * key_head_num_stride;
|
||||
scalar_t* key_cache_start_ptr =
|
||||
key_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset;
|
||||
reinterpret_cast<scalar_t*>(key_cache) +
|
||||
block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
|
||||
block_offset;
|
||||
|
||||
#pragma GCC unroll 8
|
||||
for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
|
||||
@@ -234,8 +309,9 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
|
||||
token_idx * value_token_num_stride +
|
||||
head_idx * value_head_num_stride;
|
||||
scalar_t* value_cache_start_ptr =
|
||||
value_cache + block_idx * num_blocks_stride +
|
||||
head_idx * cache_head_num_stride + block_offset * head_dim;
|
||||
reinterpret_cast<scalar_t*>(value_cache) +
|
||||
block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
|
||||
block_offset * head_dim;
|
||||
std::memcpy(value_cache_start_ptr, value_start_ptr,
|
||||
sizeof(scalar_t) * head_dim);
|
||||
}
|
||||
@@ -243,6 +319,7 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace cpu_attention
|
||||
|
||||
#endif
|
||||
|
||||
@@ -116,9 +116,9 @@ class TileGemm161 {
|
||||
} // namespace
|
||||
|
||||
// This is a general but naive implementation based on vector instructions
|
||||
template <typename scalar_t, int64_t head_dim>
|
||||
class AttentionImpl<ISA::VEC16, scalar_t, head_dim>
|
||||
: public AttentionImpl<ISA::VEC, scalar_t, head_dim> {
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::VEC16, scalar_t, head_dim, kv_cache_scalar_t>
|
||||
: public AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
|
||||
@@ -244,8 +244,8 @@ class TileGemmS390X {
|
||||
|
||||
} // namespace
|
||||
|
||||
template <typename scalar_t, int64_t head_dim>
|
||||
class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
|
||||
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
|
||||
class AttentionImpl<ISA::VXE, scalar_t, head_dim, kv_cache_scalar_t> {
|
||||
public:
|
||||
using query_t = scalar_t;
|
||||
using q_buffer_t = float;
|
||||
@@ -342,7 +342,8 @@ class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
|
||||
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 int64_t block_size, const int64_t block_size_stride,
|
||||
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
|
||||
#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) {
|
||||
|
||||
@@ -15,6 +15,9 @@ using namespace at::vec;
|
||||
|
||||
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::Half, __VA_ARGS__) \
|
||||
@@ -322,6 +325,9 @@ struct BF16Vec32 : public VectorizedRegWrapper<BF16Vec32, 4, c10::BFloat16> {
|
||||
reg.val[2] = vec8_data.reg.val[0];
|
||||
reg.val[3] = vec8_data.reg.val[0];
|
||||
};
|
||||
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : Base() {}
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : Base() {}
|
||||
};
|
||||
|
||||
struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
|
||||
|
||||
@@ -8,6 +8,9 @@
|
||||
#include <torch/all.h>
|
||||
namespace vec_op {
|
||||
|
||||
struct fp8_e4m3_tag {};
|
||||
struct fp8_e5m2_tag {};
|
||||
|
||||
#define vec_neg(a) (-(a))
|
||||
#define vec_add(a, b) ((a) + (b))
|
||||
#define vec_sub(a, b) ((a) - (b))
|
||||
@@ -241,6 +244,9 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
explicit BF16Vec32(const BF16Vec8& vec8_data)
|
||||
: reg({vec8_data.reg, vec8_data.reg, vec8_data.reg, vec8_data.reg}) {}
|
||||
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : reg{} {}
|
||||
explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : reg{} {}
|
||||
|
||||
void save(void* ptr) const { *reinterpret_cast<ss16x8x4_t*>(ptr) = reg; }
|
||||
};
|
||||
|
||||
|
||||
@@ -11,6 +11,17 @@ static_assert(false, "AVX2 must be supported for the current implementation.");
|
||||
|
||||
namespace vec_op {
|
||||
|
||||
// Tags for FP8 BF16Vec32 constructors (avoid overload collision with
|
||||
// BF16Vec32(void*)).
|
||||
// VEC path (FP8 → pseudo-FP16 layout, scale correction applied later):
|
||||
struct fp8_e4m3_tag {}; // E4M3 → pseudo-FP16; BF16 value = true_E4M3 * 2^-8
|
||||
struct fp8_e5m2_tag {}; // E5M2 → FP16 bits directly (same exponent bias=15)
|
||||
// AMX path (FP8 → unscaled BF16, no FP32 round-trip):
|
||||
// BF16 value = true_E4M3 * 2^-120 (E4M3) or true_E5M2 * 2^-112 (E5M2).
|
||||
// Exponent rebiasing is folded into k/v scales by the caller.
|
||||
struct fp8_bf16_e4m3_tag {};
|
||||
struct fp8_bf16_e5m2_tag {};
|
||||
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
|
||||
@@ -176,6 +187,50 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
(__m128i)vec8_data.reg, 2),
|
||||
(__m128i)vec8_data.reg, 3)) {}
|
||||
|
||||
// Decode 32 FP8-E4M3 bytes to pseudo-FP16 layout (stored in the BF16
|
||||
// register). Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag) {
|
||||
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
|
||||
__m512i b16 = _mm512_cvtepu8_epi16(b8);
|
||||
__m512i sign =
|
||||
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
|
||||
__m512i payload =
|
||||
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 7);
|
||||
reg = _mm512_or_si512(sign, payload);
|
||||
}
|
||||
|
||||
// Decode 32 FP8-E5M2 bytes to FP16 layout.
|
||||
// E5M2 and FP16 share the same 5-bit exponent bias (15), so FP8 byte b maps
|
||||
// directly to FP16 bits by shifting left 8 — no sign/payload reconstruction.
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag) {
|
||||
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
|
||||
reg = _mm512_slli_epi16(_mm512_cvtepu8_epi16(b8), 8);
|
||||
}
|
||||
|
||||
// Direct FP8-E4M3 → unscaled BF16 for AMX (no FP32 round-trip).
|
||||
// BF16 value = true_E4M3 * 2^-120; exponent rebiasing folded into k/v scales.
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e4m3_tag) {
|
||||
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
|
||||
__m512i b16 = _mm512_cvtepu8_epi16(b8);
|
||||
__m512i sign =
|
||||
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
|
||||
__m512i payload =
|
||||
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 4);
|
||||
reg = _mm512_or_si512(sign, payload);
|
||||
}
|
||||
|
||||
// Direct FP8-E5M2 → unscaled BF16 for AMX (no FP32 round-trip).
|
||||
// BF16 value = true_E5M2 * 2^-112; exponent rebiasing folded into k/v scales.
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e5m2_tag) {
|
||||
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
|
||||
__m512i b16 = _mm512_cvtepu8_epi16(b8);
|
||||
__m512i sign =
|
||||
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
|
||||
__m512i payload =
|
||||
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 5);
|
||||
reg = _mm512_or_si512(sign, payload);
|
||||
}
|
||||
|
||||
void save(void* ptr) const { *reinterpret_cast<__m512i*>(ptr) = reg; }
|
||||
};
|
||||
#else
|
||||
@@ -200,6 +255,77 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
_mm256_castsi128_si256((__m128i)vec8_data.reg),
|
||||
(__m128i)vec8_data.reg, 1)) {}
|
||||
|
||||
// E4M3 decode (AVX2 path) — same bit-layout trick as the AVX512 variant
|
||||
// above. Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag) {
|
||||
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
|
||||
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
|
||||
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
|
||||
__m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
|
||||
__m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
|
||||
|
||||
__m256i sign_low = _mm256_slli_epi16(
|
||||
_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)), 8);
|
||||
__m256i payload_low = _mm256_slli_epi16(
|
||||
_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)), 7);
|
||||
__m256i sign_high = _mm256_slli_epi16(
|
||||
_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)), 8);
|
||||
__m256i payload_high = _mm256_slli_epi16(
|
||||
_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)), 7);
|
||||
reg_low = _mm256_or_si256(sign_low, payload_low);
|
||||
reg_high = _mm256_or_si256(sign_high, payload_high);
|
||||
}
|
||||
|
||||
// E5M2 decode (AVX2 path) — b << 8 maps to FP16 bits; see AVX512 variant
|
||||
// above.
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag) {
|
||||
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
|
||||
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
|
||||
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
|
||||
reg_low = _mm256_slli_epi16(_mm256_cvtepu8_epi16(b8_low), 8);
|
||||
reg_high = _mm256_slli_epi16(_mm256_cvtepu8_epi16(b8_high), 8);
|
||||
}
|
||||
|
||||
// Direct FP8-E4M3 → unscaled BF16 for AMX (AVX2 path, no FP32 round-trip).
|
||||
// BF16 value = true_E4M3 * 2^-120; exponent rebiasing folded into k/v scales.
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e4m3_tag) {
|
||||
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
|
||||
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
|
||||
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
|
||||
__m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
|
||||
__m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
|
||||
reg_low = _mm256_or_si256(
|
||||
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)),
|
||||
8),
|
||||
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)),
|
||||
4));
|
||||
reg_high = _mm256_or_si256(
|
||||
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)),
|
||||
8),
|
||||
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)),
|
||||
4));
|
||||
}
|
||||
|
||||
// Direct FP8-E5M2 → unscaled BF16 for AMX (AVX2 path, no FP32 round-trip).
|
||||
// BF16 value = true_E5M2 * 2^-112; exponent rebiasing folded into k/v scales.
|
||||
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e5m2_tag) {
|
||||
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
|
||||
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
|
||||
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
|
||||
__m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
|
||||
__m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
|
||||
reg_low = _mm256_or_si256(
|
||||
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)),
|
||||
8),
|
||||
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)),
|
||||
5));
|
||||
reg_high = _mm256_or_si256(
|
||||
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)),
|
||||
8),
|
||||
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)),
|
||||
5));
|
||||
}
|
||||
|
||||
void save(void* ptr) const {
|
||||
_mm256_storeu_si256((__m256i*)ptr, reg_low);
|
||||
_mm256_storeu_si256((__m256i*)ptr + 1, reg_high);
|
||||
@@ -390,6 +516,11 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
: reg(_mm512_castsi512_ps(
|
||||
_mm512_bslli_epi128(_mm512_cvtepu16_epi32(v.reg), 2))) {}
|
||||
|
||||
explicit FP32Vec16(const BF16Vec32& v, int upper) {
|
||||
__m256i v_half_i = _mm512_extracti32x8_epi32(v.reg, upper);
|
||||
reg = _mm512_cvtph_ps(v_half_i);
|
||||
}
|
||||
|
||||
explicit FP32Vec16(const FP16Vec16& v) : reg(_mm512_cvtph_ps(v.reg)) {}
|
||||
|
||||
explicit FP32Vec16(const FP16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
|
||||
@@ -494,6 +625,14 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
explicit FP32Vec16(const FP32Vec8& data)
|
||||
: reg_low(data.reg), reg_high(data.reg) {}
|
||||
|
||||
explicit FP32Vec16(const BF16Vec32& v, int upper) {
|
||||
const __m256i& half = upper ? v.reg_high : v.reg_low;
|
||||
__m128i lo = _mm256_extractf128_si256(half, 0);
|
||||
__m128i hi = _mm256_extractf128_si256(half, 1);
|
||||
reg_low = _mm256_cvtph_ps(lo);
|
||||
reg_high = _mm256_cvtph_ps(hi);
|
||||
}
|
||||
|
||||
explicit FP32Vec16(const FP16Vec16& v) {
|
||||
__m128i low = _mm256_extractf128_si256(v.reg, 0);
|
||||
__m128i high = _mm256_extractf128_si256(v.reg, 1);
|
||||
|
||||
@@ -22,71 +22,95 @@ ISA_TYPES = {
|
||||
"VXE": 4,
|
||||
}
|
||||
|
||||
# KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
|
||||
KV_CACHE_IDX = {
|
||||
"auto": 0,
|
||||
"fp8_e4m3": 1,
|
||||
"fp8_e5m2": 2,
|
||||
}
|
||||
|
||||
# C++ type for each kv_cache index
|
||||
KV_CACHE_CPP_TYPES = {
|
||||
"auto": "scalar_t",
|
||||
"fp8_e4m3": "c10::Float8_e4m3fn",
|
||||
"fp8_e5m2": "c10::Float8_e5m2",
|
||||
}
|
||||
|
||||
# ISAs supported for head_dims divisible by 32
|
||||
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE"]
|
||||
|
||||
# ISAs supported for head_dims divisible by 16 only
|
||||
ISA_FOR_16 = ["VEC16"]
|
||||
|
||||
# ISAs that support FP8 KV cache (x86 AVX2/AVX-512 required)
|
||||
ISA_FOR_FP8 = ["AMX", "VEC"]
|
||||
|
||||
def encode_params(head_dim: int, isa_type: str) -> int:
|
||||
"""Encode head_dim and ISA type into a single int64_t."""
|
||||
|
||||
def encode_params(head_dim: int, isa_type: str, kv_cache: str = "auto") -> int:
|
||||
"""Encode head_dim, ISA type, and KV cache type into a single int64_t."""
|
||||
isa_val = ISA_TYPES[isa_type]
|
||||
# Encoding: (head_dim << 8) | isa_type
|
||||
# This allows head_dim up to 2^56 - 1 and 256 ISA types
|
||||
return (head_dim << 8) | isa_val
|
||||
kv_val = KV_CACHE_IDX[kv_cache]
|
||||
# Encoding: (head_dim << 16) | (kv_cache_idx << 8) | isa_type
|
||||
# This allows head_dim up to 2^48 - 1, 256 KV cache types, and 256 ISA types
|
||||
return (head_dim << 16) | (kv_val << 8) | isa_val
|
||||
|
||||
|
||||
def generate_cases_for_isa_group(isa_list: list[str]) -> str:
|
||||
def _make_case(
|
||||
head_dim: int, isa: str, kv_cache: str = "auto", isa_override: str | None = None
|
||||
) -> str:
|
||||
"""Generate a single switch case line."""
|
||||
encoded = encode_params(head_dim, isa, kv_cache)
|
||||
actual_isa = isa_override if isa_override else isa
|
||||
cpp_type = KV_CACHE_CPP_TYPES[kv_cache]
|
||||
attn_impl = (
|
||||
f"cpu_attention::AttentionImpl<"
|
||||
f"cpu_attention::ISA::{actual_isa}, \\\n"
|
||||
f" "
|
||||
f"scalar_t, head_dim, {cpp_type}>"
|
||||
)
|
||||
comment = (
|
||||
f"head_dim={head_dim}, isa={isa}"
|
||||
if kv_cache == "auto"
|
||||
else f"head_dim={head_dim}, isa={isa}, kv_cache={kv_cache}"
|
||||
)
|
||||
return (
|
||||
f""" case {encoded}LL: {{ """
|
||||
f"""/* {comment} */ \\"""
|
||||
f"""
|
||||
constexpr size_t head_dim = {head_dim}; \\"""
|
||||
f"""
|
||||
using attn_impl = {attn_impl}; \\"""
|
||||
f"""
|
||||
return __VA_ARGS__(); \\"""
|
||||
f"""
|
||||
}} \\"""
|
||||
)
|
||||
|
||||
|
||||
def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False) -> str:
|
||||
"""Generate switch cases for a specific ISA group."""
|
||||
cases = []
|
||||
|
||||
# Generate cases for head_dims divisible by 32
|
||||
# Non-FP8 cases for head_dims divisible by 32
|
||||
for head_dim in HEAD_DIMS_32:
|
||||
for isa in isa_list:
|
||||
if isa not in ISA_FOR_32:
|
||||
continue
|
||||
encoded = encode_params(head_dim, isa)
|
||||
case_str = (
|
||||
f""" case {encoded}LL: {{ """
|
||||
f"""/* head_dim={head_dim}, isa={isa} */ \\"""
|
||||
f"""
|
||||
constexpr size_t head_dim = {head_dim}; \\"""
|
||||
f"""
|
||||
using attn_impl = cpu_attention::AttentionImpl<"""
|
||||
f"""cpu_attention::ISA::{isa}, \\"""
|
||||
f"""
|
||||
"""
|
||||
f"""scalar_t, head_dim>; \\"""
|
||||
f"""
|
||||
return __VA_ARGS__(); \\"""
|
||||
f"""
|
||||
}} \\"""
|
||||
)
|
||||
cases.append(case_str)
|
||||
cases.append(_make_case(head_dim, isa, "auto"))
|
||||
|
||||
# Generate cases for head_dims divisible by 16 only
|
||||
# Non-FP8 cases for head_dims divisible by 16 only
|
||||
for head_dim in HEAD_DIMS_16:
|
||||
for isa in isa_list:
|
||||
encoded = encode_params(head_dim, isa)
|
||||
case_str = (
|
||||
f""" case {encoded}LL: {{ """
|
||||
f"""/* head_dim={head_dim}, isa={isa} """
|
||||
f"""(using VEC16) */ \\"""
|
||||
f"""
|
||||
constexpr size_t head_dim = {head_dim}; \\"""
|
||||
f"""
|
||||
using attn_impl = cpu_attention::AttentionImpl<"""
|
||||
f"""cpu_attention::ISA::VEC16, \\"""
|
||||
f"""
|
||||
"""
|
||||
f"""scalar_t, head_dim>; \\"""
|
||||
f"""
|
||||
return __VA_ARGS__(); \\"""
|
||||
f"""
|
||||
}} \\"""
|
||||
)
|
||||
cases.append(case_str)
|
||||
cases.append(_make_case(head_dim, isa, "auto", isa_override="VEC16"))
|
||||
|
||||
# FP8 cases: only AMX and VEC, only head_dims divisible by 32
|
||||
if include_fp8:
|
||||
for fp8_type in ("fp8_e4m3", "fp8_e5m2"):
|
||||
for head_dim in HEAD_DIMS_32:
|
||||
for isa in isa_list:
|
||||
if isa not in ISA_FOR_FP8:
|
||||
continue
|
||||
cases.append(_make_case(head_dim, isa, fp8_type))
|
||||
|
||||
return "\n".join(cases)
|
||||
|
||||
@@ -94,8 +118,9 @@ def generate_cases_for_isa_group(isa_list: list[str]) -> str:
|
||||
def generate_helper_function() -> str:
|
||||
"""Generate helper function to encode parameters."""
|
||||
return """
|
||||
inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa) {
|
||||
return (head_dim << 8) | static_cast<int64_t>(isa);
|
||||
inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa,
|
||||
int64_t kv_cache_idx = 0) {
|
||||
return (head_dim << 16) | (kv_cache_idx << 8) | static_cast<int64_t>(isa);
|
||||
}
|
||||
"""
|
||||
|
||||
@@ -129,87 +154,78 @@ def generate_header_file() -> str:
|
||||
|
||||
# Generate dispatch macro with conditional compilation for different ISA sets
|
||||
header += """
|
||||
// Dispatch macro using encoded parameters
|
||||
// Dispatch macro using encoded parameters.
|
||||
// KV_CACHE_IDX: Fp8KVCacheDataType enum value (kAuto=0, kFp8E4M3=1, kFp8E5M2=2).
|
||||
// FP8 cases (kv_cache_idx != 0) are generated on x86 platforms with AVX2 or
|
||||
// AVX-512: BF16Vec32 FP8 constructors have both AVX-512 and AVX2 implementations
|
||||
// in cpu_types_x86.hpp. Non-x86 platforms (#else fallback) have fp8=False.
|
||||
"""
|
||||
|
||||
# x86_64 with AMX
|
||||
header += """#if defined(CPU_CAPABILITY_AMXBF16)
|
||||
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
|
||||
[&] { \\
|
||||
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
|
||||
switch (encoded_params) { \\
|
||||
"""
|
||||
header += generate_cases_for_isa_group(["AMX", "VEC", "VEC16"])
|
||||
header += """
|
||||
default: { \\
|
||||
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
|
||||
std::to_string(HEAD_DIM) + " isa=" + \\
|
||||
std::to_string(static_cast<int>(ISA_TYPE))); \\
|
||||
} \\
|
||||
} \\
|
||||
}()
|
||||
def _macro_block(guard: str, isa_list: list[str], fp8: bool) -> str:
|
||||
"""Return one CPU_ATTN_DISPATCH macro block for a given guard."""
|
||||
enc = (
|
||||
" int64_t encoded_params = encode_cpu_attn_params("
|
||||
"HEAD_DIM, ISA_TYPE, KV_CACHE_IDX); \\"
|
||||
)
|
||||
cases = generate_cases_for_isa_group(isa_list, include_fp8=fp8)
|
||||
tail = (
|
||||
"\n"
|
||||
" default: { \\\n"
|
||||
" TORCH_CHECK(false, "
|
||||
'"Unsupported CPU attention configuration: head_dim=" + \\\n'
|
||||
' std::to_string(HEAD_DIM) + " isa=" + \\\n'
|
||||
" std::to_string(static_cast<int>(ISA_TYPE))"
|
||||
" + \\\n"
|
||||
' " kv_cache_idx=" + '
|
||||
"std::to_string(KV_CACHE_IDX)); \\\n"
|
||||
" } \\\n"
|
||||
" } \\\n"
|
||||
" }()\n\n"
|
||||
)
|
||||
return (
|
||||
f"{guard}\n"
|
||||
"#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, KV_CACHE_IDX, ...) \\\n"
|
||||
" [&] { \\\n"
|
||||
f"{enc}\n"
|
||||
" switch (encoded_params) { \\\n"
|
||||
f"{cases}"
|
||||
f"{tail}"
|
||||
)
|
||||
|
||||
"""
|
||||
|
||||
# ARM64 with NEON
|
||||
header += """#elif defined(__aarch64__)
|
||||
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
|
||||
[&] { \\
|
||||
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
|
||||
switch (encoded_params) { \\
|
||||
"""
|
||||
header += generate_cases_for_isa_group(["NEON", "VEC", "VEC16"])
|
||||
header += """
|
||||
default: { \\
|
||||
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
|
||||
std::to_string(HEAD_DIM) + " isa=" + \\
|
||||
std::to_string(static_cast<int>(ISA_TYPE))); \\
|
||||
} \\
|
||||
} \\
|
||||
}()
|
||||
|
||||
"""
|
||||
|
||||
# s390x with VXE
|
||||
header += """#elif defined(__s390x__)
|
||||
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
|
||||
[&] { \\
|
||||
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
|
||||
switch (encoded_params) { \\
|
||||
"""
|
||||
header += generate_cases_for_isa_group(["VXE", "VEC", "VEC16"])
|
||||
header += """
|
||||
default: { \\
|
||||
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
|
||||
std::to_string(HEAD_DIM) + " isa=" + \\
|
||||
std::to_string(static_cast<int>(ISA_TYPE))); \\
|
||||
} \\
|
||||
} \\
|
||||
}()
|
||||
|
||||
"""
|
||||
|
||||
# Fallback: VEC and VEC16 only
|
||||
header += """#else
|
||||
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
|
||||
[&] { \\
|
||||
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
|
||||
switch (encoded_params) { \\
|
||||
"""
|
||||
header += generate_cases_for_isa_group(["VEC", "VEC16"])
|
||||
header += """
|
||||
default: { \\
|
||||
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
|
||||
std::to_string(HEAD_DIM) + " isa=" + \\
|
||||
std::to_string(static_cast<int>(ISA_TYPE))); \\
|
||||
} \\
|
||||
} \\
|
||||
}()
|
||||
|
||||
#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */
|
||||
|
||||
#endif // CPU_ATTN_DISPATCH_GENERATED_H
|
||||
"""
|
||||
header += _macro_block(
|
||||
"#if defined(CPU_CAPABILITY_AMXBF16)",
|
||||
["AMX", "VEC", "VEC16"],
|
||||
fp8=True,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__aarch64__)",
|
||||
["NEON", "VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__s390x__)",
|
||||
["VXE", "VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__AVX512F__)",
|
||||
["VEC", "VEC16"],
|
||||
fp8=True,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#elif defined(__AVX2__)",
|
||||
["VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += _macro_block(
|
||||
"#else",
|
||||
["VEC", "VEC16"],
|
||||
fp8=False,
|
||||
)
|
||||
header += (
|
||||
"#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */\n\n"
|
||||
"#endif // CPU_ATTN_DISPATCH_GENERATED_H\n"
|
||||
)
|
||||
|
||||
return header
|
||||
|
||||
|
||||
@@ -101,7 +101,9 @@ void cpu_attn_reshape_and_cache(const torch::Tensor& key,
|
||||
torch::Tensor& key_cache,
|
||||
torch::Tensor& value_cache,
|
||||
const torch::Tensor& slot_mapping,
|
||||
const std::string& isa);
|
||||
const std::string& isa, const double k_scale,
|
||||
const double v_scale,
|
||||
const std::string& kv_cache_dtype);
|
||||
|
||||
void cpu_attention_with_kv_cache(
|
||||
const torch::Tensor& query, const torch::Tensor& key_cache,
|
||||
@@ -112,7 +114,8 @@ void cpu_attention_with_kv_cache(
|
||||
const int64_t sliding_window_left, const int64_t sliding_window_right,
|
||||
const torch::Tensor& block_table, const double softcap,
|
||||
const torch::Tensor& scheduler_metadata,
|
||||
const std::optional<torch::Tensor>& s_aux);
|
||||
const std::optional<torch::Tensor>& s_aux, const double k_scale,
|
||||
const double v_scale, const std::string& kv_cache_dtype);
|
||||
|
||||
// Note: just for avoiding importing errors
|
||||
void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); }
|
||||
@@ -384,15 +387,18 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
&get_scheduler_metadata);
|
||||
ops.def(
|
||||
"cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
|
||||
"key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str "
|
||||
"isa) -> ()",
|
||||
"key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str isa, "
|
||||
"float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
|
||||
"()",
|
||||
&cpu_attn_reshape_and_cache);
|
||||
ops.def(
|
||||
"cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor "
|
||||
"value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor "
|
||||
"seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt "
|
||||
"sliding_window_left, SymInt sliding_window_right, Tensor block_table, "
|
||||
"float softcap, Tensor scheduler_metadata, Tensor? s_aux) -> ()",
|
||||
"float softcap, Tensor scheduler_metadata, Tensor? s_aux, "
|
||||
"float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
|
||||
"()",
|
||||
&cpu_attention_with_kv_cache);
|
||||
|
||||
// placeholders
|
||||
|
||||
@@ -96,44 +96,14 @@ struct enable_sm90_or_later : Kernel {
|
||||
};
|
||||
|
||||
template <typename Kernel>
|
||||
struct enable_sm90_only : Kernel {
|
||||
struct enable_sm100_to_sm120 : Kernel {
|
||||
template <typename... Args>
|
||||
CUTLASS_DEVICE void operator()(Args&&... args) {
|
||||
#if defined __CUDA_ARCH__
|
||||
#if __CUDA_ARCH__ == 900
|
||||
#if (__CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1200)
|
||||
Kernel::operator()(std::forward<Args>(args)...);
|
||||
#else
|
||||
printf("This kernel only supports sm90.\n");
|
||||
asm("trap;");
|
||||
#endif
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Kernel>
|
||||
struct enable_sm100f_only : Kernel {
|
||||
template <typename... Args>
|
||||
CUTLASS_DEVICE void operator()(Args&&... args) {
|
||||
#if defined __CUDA_ARCH__
|
||||
#if __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030
|
||||
Kernel::operator()(std::forward<Args>(args)...);
|
||||
#else
|
||||
printf("This kernel only supports sm100f.\n");
|
||||
asm("trap;");
|
||||
#endif
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
template <typename Kernel>
|
||||
struct enable_sm100a_only : Kernel {
|
||||
template <typename... Args>
|
||||
CUTLASS_DEVICE void operator()(Args&&... args) {
|
||||
#if defined __CUDA_ARCH__
|
||||
#if __CUDA_ARCH__ == 1000
|
||||
Kernel::operator()(std::forward<Args>(args)...);
|
||||
#else
|
||||
printf("This kernel only supports sm100a.\n");
|
||||
printf("This kernel only supports sm[100, 120).\n");
|
||||
asm("trap;");
|
||||
#endif
|
||||
#endif
|
||||
@@ -148,7 +118,7 @@ struct enable_sm120_only : Kernel {
|
||||
#if __CUDA_ARCH__ == 1200
|
||||
Kernel::operator()(std::forward<Args>(args)...);
|
||||
#else
|
||||
printf("This kernel only supports sm120.\n");
|
||||
printf("This kernel only supports sm120a.\n");
|
||||
asm("trap;");
|
||||
#endif
|
||||
#endif
|
||||
@@ -160,8 +130,13 @@ template <typename Kernel>
|
||||
struct enable_sm120_family : Kernel {
|
||||
template <typename... Args>
|
||||
CUTLASS_DEVICE void operator()(Args&&... args) {
|
||||
#if defined __CUDA_ARCH__ && (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
|
||||
#if defined __CUDA_ARCH__
|
||||
#if (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
|
||||
Kernel::operator()(std::forward<Args>(args)...);
|
||||
#else
|
||||
printf("This kernel only supports sm120f.\n");
|
||||
asm("trap;");
|
||||
#endif
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
@@ -141,7 +141,7 @@ struct cutlass_3x_gemm_sm100 {
|
||||
sizeof(typename CollectiveEpilogue::SharedStorage))>,
|
||||
KernelSchedule>::CollectiveOp;
|
||||
|
||||
using GemmKernel = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
|
||||
using GemmKernel = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
|
||||
};
|
||||
|
||||
|
||||
+1
-1
@@ -125,7 +125,7 @@ struct cutlass_3x_gemm_fp8_blockwise {
|
||||
MainloopScheduler
|
||||
>::CollectiveOp>;
|
||||
|
||||
using KernelType = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
|
||||
using KernelType = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue>>;
|
||||
|
||||
struct GemmKernel : public KernelType {};
|
||||
|
||||
@@ -92,7 +92,7 @@ struct cutlass_3x_gemm_sm100_fp8 {
|
||||
// -----------------------------------------------------------
|
||||
// Kernel definition
|
||||
// -----------------------------------------------------------
|
||||
using GemmKernel = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
|
||||
using GemmKernel = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
|
||||
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
|
||||
};
|
||||
|
||||
|
||||
@@ -7,23 +7,23 @@
|
||||
|
||||
namespace vllm {
|
||||
|
||||
template <typename scalar_t, bool IS_NEOX>
|
||||
template <typename scalar_t, typename cache_t, bool IS_NEOX>
|
||||
inline __device__ void apply_token_rotary_embedding(
|
||||
scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
|
||||
const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
|
||||
scalar_t* __restrict__ arr, const cache_t* __restrict__ cos_ptr,
|
||||
const cache_t* __restrict__ sin_ptr, int rot_offset, int embed_dim,
|
||||
const bool inverse) {
|
||||
int x_index, y_index;
|
||||
float cos_f, sin_f;
|
||||
if (IS_NEOX) {
|
||||
x_index = rot_offset;
|
||||
y_index = embed_dim + rot_offset;
|
||||
cos_f = VLLM_LDG(cos_ptr + x_index);
|
||||
sin_f = VLLM_LDG(sin_ptr + x_index);
|
||||
cos_f = static_cast<float>(VLLM_LDG(cos_ptr + x_index));
|
||||
sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index));
|
||||
} else {
|
||||
x_index = 2 * rot_offset;
|
||||
y_index = 2 * rot_offset + 1;
|
||||
cos_f = VLLM_LDG(cos_ptr + x_index / 2);
|
||||
sin_f = VLLM_LDG(sin_ptr + x_index / 2);
|
||||
cos_f = static_cast<float>(VLLM_LDG(cos_ptr + x_index / 2));
|
||||
sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index / 2));
|
||||
}
|
||||
if (inverse) {
|
||||
sin_f = -sin_f;
|
||||
@@ -34,7 +34,7 @@ inline __device__ void apply_token_rotary_embedding(
|
||||
arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
|
||||
}
|
||||
|
||||
template <typename scalar_t, bool IS_NEOX>
|
||||
template <typename scalar_t, typename cache_t, bool IS_NEOX>
|
||||
inline __device__ void apply_rotary_embedding(
|
||||
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
|
||||
// head_size] or [num_tokens, num_heads,
|
||||
@@ -43,14 +43,14 @@ inline __device__ void apply_rotary_embedding(
|
||||
// [batch_size, seq_len, num_kv_heads,
|
||||
// head_size] or [num_tokens, num_kv_heads,
|
||||
// head_size]
|
||||
const float* cache_ptr, const int head_size, const int num_heads,
|
||||
const cache_t* cache_ptr, const int head_size, const int num_heads,
|
||||
const int num_kv_heads, const int rot_dim, const int token_idx,
|
||||
const int64_t query_stride, const int64_t key_stride,
|
||||
const int64_t head_stride, const int64_t rope_dim_offset,
|
||||
const bool inverse) {
|
||||
const int embed_dim = rot_dim / 2;
|
||||
const float* cos_ptr = cache_ptr;
|
||||
const float* sin_ptr = cache_ptr + embed_dim;
|
||||
const cache_t* cos_ptr = cache_ptr;
|
||||
const cache_t* sin_ptr = cache_ptr + embed_dim;
|
||||
|
||||
const int nq = num_heads * embed_dim;
|
||||
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
|
||||
@@ -58,7 +58,7 @@ inline __device__ void apply_rotary_embedding(
|
||||
const int64_t token_head =
|
||||
token_idx * query_stride + head_idx * head_stride + rope_dim_offset;
|
||||
const int rot_offset = i % embed_dim;
|
||||
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
|
||||
apply_token_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
|
||||
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
|
||||
}
|
||||
|
||||
@@ -69,13 +69,13 @@ inline __device__ void apply_rotary_embedding(
|
||||
const int64_t token_head =
|
||||
token_idx * key_stride + head_idx * head_stride + rope_dim_offset;
|
||||
const int rot_offset = i % embed_dim;
|
||||
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
|
||||
apply_token_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
|
||||
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t, bool IS_NEOX>
|
||||
template <typename scalar_t, typename cache_t, bool IS_NEOX>
|
||||
__global__ void rotary_embedding_kernel(
|
||||
const int64_t* __restrict__ positions, // [batch_size, seq_len] or
|
||||
// [num_tokens]
|
||||
@@ -86,15 +86,15 @@ __global__ void rotary_embedding_kernel(
|
||||
// [batch_size, seq_len, num_kv_heads,
|
||||
// head_size] or [num_tokens, num_kv_heads,
|
||||
// head_size]
|
||||
const float* __restrict__ cos_sin_cache, // [max_position, rot_dim] fp32
|
||||
const cache_t* __restrict__ cos_sin_cache, // [max_position, rot_dim]
|
||||
const int rot_dim, const int64_t query_stride, const int64_t key_stride,
|
||||
const int64_t head_stride, const int num_heads, const int num_kv_heads,
|
||||
const int head_size, const int64_t rope_dim_offset, const bool inverse) {
|
||||
const int token_idx = blockIdx.x;
|
||||
int64_t pos = positions[token_idx];
|
||||
const float* cache_ptr = cos_sin_cache + pos * rot_dim;
|
||||
const cache_t* cache_ptr = cos_sin_cache + pos * rot_dim;
|
||||
|
||||
apply_rotary_embedding<scalar_t, IS_NEOX>(
|
||||
apply_rotary_embedding<scalar_t, cache_t, IS_NEOX>(
|
||||
query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
|
||||
token_idx, query_stride, key_stride, head_stride, rope_dim_offset,
|
||||
inverse);
|
||||
@@ -168,23 +168,28 @@ void rotary_embedding(
|
||||
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
auto cache_f32 = cos_sin_cache.to(torch::kFloat32);
|
||||
VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
|
||||
if (is_neox) {
|
||||
vllm::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
|
||||
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
|
||||
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
|
||||
cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
|
||||
head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
|
||||
inverse);
|
||||
} else {
|
||||
vllm::rotary_embedding_kernel<scalar_t, false>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
|
||||
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
|
||||
cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
|
||||
head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
|
||||
inverse);
|
||||
}
|
||||
using query_t = scalar_t;
|
||||
VLLM_DISPATCH_FLOATING_TYPES(
|
||||
cos_sin_cache.scalar_type(), "rotary_embedding_cache", [&] {
|
||||
using cache_t = scalar_t;
|
||||
if (is_neox) {
|
||||
vllm::rotary_embedding_kernel<query_t, cache_t, true>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
|
||||
key.has_value() ? key->data_ptr<query_t>() : nullptr,
|
||||
cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
|
||||
key_stride, head_stride, num_heads, num_kv_heads, head_size,
|
||||
rope_dim_offset, inverse);
|
||||
} else {
|
||||
vllm::rotary_embedding_kernel<query_t, cache_t, false>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
positions.data_ptr<int64_t>(), query.data_ptr<query_t>(),
|
||||
key.has_value() ? key->data_ptr<query_t>() : nullptr,
|
||||
cos_sin_cache.data_ptr<cache_t>(), rot_dim, query_stride,
|
||||
key_stride, head_stride, num_heads, num_kv_heads, head_size,
|
||||
rope_dim_offset, inverse);
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
+59
-4
@@ -82,18 +82,73 @@ void launch_persistent_topk(const torch::Tensor& logits,
|
||||
size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t);
|
||||
if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium;
|
||||
|
||||
// Query occupancy for the instantiation that will actually launch;
|
||||
// overestimating it deadlocks the cooperative barrier.
|
||||
int occupancy = 1;
|
||||
cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
&occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
|
||||
smem_size);
|
||||
cudaError_t occ_err = cudaSuccess;
|
||||
if (vec_size == 4) {
|
||||
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
&occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
|
||||
smem_size);
|
||||
} else if (vec_size == 2) {
|
||||
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
&occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock,
|
||||
smem_size);
|
||||
} else {
|
||||
occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
&occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock,
|
||||
smem_size);
|
||||
}
|
||||
TORCH_CHECK(occ_err == cudaSuccess,
|
||||
"persistent_topk occupancy query failed: ",
|
||||
cudaGetErrorString(occ_err));
|
||||
if (occupancy < 1) occupancy = 1;
|
||||
|
||||
uint32_t max_resident_ctas = static_cast<uint32_t>(num_sms) * occupancy;
|
||||
// The cooperative spin-wait barrier only runs when at least one row hits
|
||||
// the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs
|
||||
// early-exit, so oversubscription can't deadlock and headroom is wasted.
|
||||
const bool needs_cooperative =
|
||||
static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD;
|
||||
|
||||
const uint32_t hw_resident_cap =
|
||||
static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy);
|
||||
uint32_t max_resident_ctas = hw_resident_cap;
|
||||
if (needs_cooperative) {
|
||||
// Reserve one CTA per SM when occupancy allows; fall back to a single
|
||||
// CTA when occupancy == 1 (the most deadlock-prone case — any straggler
|
||||
// kernel that takes the only slot on one SM hangs the barrier). Never
|
||||
// drop below one full group's worth.
|
||||
uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u;
|
||||
if (max_resident_ctas >= headroom + ctas_per_group) {
|
||||
max_resident_ctas -= headroom;
|
||||
}
|
||||
}
|
||||
uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group,
|
||||
static_cast<uint32_t>(num_rows));
|
||||
if (num_groups == 0) num_groups = 1;
|
||||
uint32_t total_ctas = num_groups * ctas_per_group;
|
||||
|
||||
// If the cooperative launch wouldn't fit, fall back to FilteredTopK
|
||||
// instead of deadlocking. Only relevant when needs_cooperative.
|
||||
if (needs_cooperative && total_ctas > hw_resident_cap) {
|
||||
TORCH_CHECK(max_smem_per_block >= 128 * 1024,
|
||||
"persistent_topk would oversubscribe and the FilteredTopK "
|
||||
"fallback requires >=128KB smem per block (have ",
|
||||
max_smem_per_block, "). total_ctas=", total_ctas,
|
||||
" > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK,
|
||||
", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group,
|
||||
", smem=", smem_size, ").");
|
||||
cudaError_t status =
|
||||
vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>(
|
||||
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
|
||||
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
|
||||
static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride),
|
||||
stream);
|
||||
TORCH_CHECK(status == cudaSuccess,
|
||||
"FilteredTopK fallback failed: ", cudaGetErrorString(status));
|
||||
return;
|
||||
}
|
||||
|
||||
size_t state_bytes = num_groups * sizeof(P::RadixRowState);
|
||||
TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes),
|
||||
"workspace too small, need ", state_bytes, " bytes");
|
||||
|
||||
+27
-4
@@ -540,7 +540,9 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
|
||||
libcurand-dev-${CUDA_VERSION_DASH} \
|
||||
libcublas-dev-${CUDA_VERSION_DASH} \
|
||||
# Required by fastsafetensors (fixes #20384)
|
||||
libnuma-dev && \
|
||||
libnuma-dev \
|
||||
# numactl CLI for NUMA binding at runtime
|
||||
numactl && \
|
||||
# Fixes nccl_allocator requiring nccl.h at runtime
|
||||
# https://github.com/vllm-project/vllm/blob/1336a1ea244fa8bfd7e72751cabbdb5b68a0c11a/vllm/distributed/device_communicators/pynccl_allocator.py#L22
|
||||
# NCCL packages don't use the cuda-MAJOR-MINOR naming convention,
|
||||
@@ -583,9 +585,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
ARG FLASHINFER_VERSION=0.6.8.post1
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
|
||||
&& flashinfer show-config \
|
||||
&& flashinfer download-cubin
|
||||
--extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
# ============================================================
|
||||
# OPENAI API SERVER DEPENDENCIES
|
||||
@@ -667,6 +667,13 @@ RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm
|
||||
uv pip install --system ep_kernels/dist/*.whl --verbose \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
# Download FlashInfer precompiled cubins AFTER all pip installs are done.
|
||||
# This must run after the vLLM wheel and EP kernels installs above, because
|
||||
# those can reinstall/touch flashinfer packages. Downloading cubins earlier
|
||||
# (in the flashinfer-jit-cache layer) causes ~2.5 GB of layer duplication
|
||||
# when a later pip install overwrites flashinfer package files.
|
||||
RUN flashinfer show-config && flashinfer download-cubin
|
||||
|
||||
# CUDA image changed from /usr/local/nvidia to /usr/local/cuda in 12.8 but will
|
||||
# return to /usr/local/nvidia in 13.0 to allow container providers to mount drivers
|
||||
# consistently from the host (see https://github.com/vllm-project/vllm/issues/18859).
|
||||
@@ -756,6 +763,10 @@ FROM vllm-base AS vllm-openai-base
|
||||
ARG TARGETPLATFORM
|
||||
ARG INSTALL_KV_CONNECTORS=false
|
||||
ARG CUDA_VERSION
|
||||
ARG VLLM_BUILD_COMMIT
|
||||
ARG VLLM_BUILD_PIPELINE
|
||||
ARG VLLM_BUILD_URL
|
||||
ARG VLLM_IMAGE_TAG
|
||||
|
||||
ARG PIP_INDEX_URL UV_INDEX_URL
|
||||
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
|
||||
@@ -792,6 +803,18 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
fi
|
||||
|
||||
ENV VLLM_USAGE_SOURCE production-docker-image
|
||||
ENV VLLM_BUILD_COMMIT=${VLLM_BUILD_COMMIT:-unknown} \
|
||||
VLLM_BUILD_PIPELINE=${VLLM_BUILD_PIPELINE:-local} \
|
||||
VLLM_BUILD_URL=${VLLM_BUILD_URL:-} \
|
||||
VLLM_IMAGE_TAG=${VLLM_IMAGE_TAG:-local/vllm-openai:dev}
|
||||
LABEL org.opencontainers.image.source="https://github.com/vllm-project/vllm" \
|
||||
org.opencontainers.image.revision="${VLLM_BUILD_COMMIT}" \
|
||||
org.opencontainers.image.version="${VLLM_IMAGE_TAG}" \
|
||||
org.opencontainers.image.url="${VLLM_BUILD_URL}" \
|
||||
ai.vllm.build.commit="${VLLM_BUILD_COMMIT}" \
|
||||
ai.vllm.build.pipeline="${VLLM_BUILD_PIPELINE}" \
|
||||
ai.vllm.build.url="${VLLM_BUILD_URL}" \
|
||||
ai.vllm.image.tag="${VLLM_IMAGE_TAG}"
|
||||
|
||||
# define sagemaker first, so it is not default from `docker build`
|
||||
FROM vllm-openai-base AS vllm-sagemaker
|
||||
|
||||
@@ -192,6 +192,7 @@ ADD ./tests/ ./tests/
|
||||
ADD ./examples/ ./examples/
|
||||
ADD ./benchmarks/ ./benchmarks/
|
||||
ADD ./vllm/collect_env.py .
|
||||
ADD ./docker/ ./docker/
|
||||
ADD ./.buildkite/ ./.buildkite/
|
||||
|
||||
# install development dependencies (for testing)
|
||||
|
||||
+15
-8
@@ -200,9 +200,9 @@ RUN cd /opt/rixl && \
|
||||
|
||||
# DeepEP build stage
|
||||
FROM base AS build_deep
|
||||
ARG ROCSHMEM_BRANCH="ba0bf0f3"
|
||||
ARG ROCSHMEM_BRANCH="f0acb0c6"
|
||||
ARG ROCSHMEM_REPO="https://github.com/ROCm/rocm-systems.git"
|
||||
ARG DEEPEP_BRANCH="5d90af8b"
|
||||
ARG DEEPEP_BRANCH="a9ea9774"
|
||||
ARG DEEPEP_REPO="https://github.com/ROCm/DeepEP.git"
|
||||
ARG DEEPEP_NIC="cx7"
|
||||
ARG DEEPEP_ROCM_ARCH="gfx942;gfx950"
|
||||
@@ -213,18 +213,15 @@ RUN git clone ${ROCSHMEM_REPO} \
|
||||
&& git checkout ${ROCSHMEM_BRANCH} \
|
||||
&& mkdir -p projects/rocshmem/build \
|
||||
&& cd projects/rocshmem/build \
|
||||
&& bash ../scripts/build_configs/all_backends \
|
||||
-DCMAKE_INSTALL_PREFIX="${ROCSHMEM_DIR}" \
|
||||
-DROCM_PATH=/opt/rocm \
|
||||
-DGPU_TARGETS="${DEEPEP_ROCM_ARCH}" \
|
||||
-DUSE_EXTERNAL_MPI=OFF
|
||||
&& INSTALL_PREFIX=${ROCSHMEM_DIR} \
|
||||
../scripts/build_configs/all_backends -DUSE_EXTERNAL_MPI=OFF
|
||||
|
||||
# Build DeepEP wheel.
|
||||
# DeepEP looks for rocshmem at ROCSHMEM_DIR.
|
||||
RUN git clone ${DEEPEP_REPO} \
|
||||
&& cd DeepEP \
|
||||
&& git checkout ${DEEPEP_BRANCH} \
|
||||
&& python3 setup.py --variant rocm --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
|
||||
&& python3 setup.py --variant rocm --rocm-explicit-ctx --nic ${DEEPEP_NIC} bdist_wheel --dist-dir=/app/deep_install
|
||||
|
||||
# MoRI runtime dependencies live in Dockerfile.rocm so NIC backend changes do
|
||||
# not force users to rebuild the long-lived Dockerfile.rocm_base image.
|
||||
@@ -388,6 +385,16 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
|
||||
# above are not available once that RUN step completes.
|
||||
COPY --from=export_vllm /*.whl /opt/vllm-wheels/
|
||||
|
||||
# Update rdma-core to support latest rocshmem
|
||||
ARG DEEPEP_NIC
|
||||
RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
|
||||
git clone --branch v62.0 --depth 1 https://github.com/linux-rdma/rdma-core.git /tmp/rdma-core && \
|
||||
cd /tmp/rdma-core && \
|
||||
mkdir -p build && cd build && \
|
||||
cmake -GNinja -DCMAKE_INSTALL_PREFIX=/usr -DNO_MAN_PAGES=1 .. && \
|
||||
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
|
||||
fi
|
||||
|
||||
# Install RIXL wheel
|
||||
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
|
||||
uv pip install --system /rixl_install/*.whl
|
||||
|
||||
+14
-5
@@ -5,9 +5,6 @@ WORKDIR /workspace/
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
|
||||
|
||||
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
|
||||
|
||||
RUN apt clean && apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --fix-missing \
|
||||
curl \
|
||||
@@ -26,8 +23,20 @@ RUN apt clean && apt-get update -y && \
|
||||
python3.12-dev \
|
||||
python3-pip
|
||||
|
||||
RUN apt update && apt upgrade -y && \
|
||||
apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
|
||||
# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
|
||||
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
|
||||
printf '%s\n' \
|
||||
'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
|
||||
'Pin: version 2025.3*' \
|
||||
'Pin-Priority: 1001' \
|
||||
> /etc/apt/preferences.d/oneapi-2025.3.pref && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
intel-oneapi-compiler-dpcpp-cpp-2025.3 \
|
||||
intel-oneapi-mkl-devel-2025.3 \
|
||||
intel-oneapi-dnnl-devel-2025.3 && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install UMD
|
||||
RUN mkdir neo && \
|
||||
|
||||
+28
-2
@@ -27,6 +27,22 @@ variable "COMMIT" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
variable "VLLM_BUILD_COMMIT" {
|
||||
default = "unknown"
|
||||
}
|
||||
|
||||
variable "VLLM_BUILD_PIPELINE" {
|
||||
default = "local"
|
||||
}
|
||||
|
||||
variable "VLLM_BUILD_URL" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
variable "VLLM_IMAGE_TAG" {
|
||||
default = "local/vllm-openai:dev"
|
||||
}
|
||||
|
||||
# Groups
|
||||
|
||||
group "default" {
|
||||
@@ -46,6 +62,10 @@ target "_common" {
|
||||
max_jobs = MAX_JOBS
|
||||
nvcc_threads = NVCC_THREADS
|
||||
torch_cuda_arch_list = TORCH_CUDA_ARCH_LIST
|
||||
VLLM_BUILD_COMMIT = VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")
|
||||
VLLM_BUILD_PIPELINE = VLLM_BUILD_PIPELINE
|
||||
VLLM_BUILD_URL = VLLM_BUILD_URL
|
||||
VLLM_IMAGE_TAG = VLLM_IMAGE_TAG
|
||||
}
|
||||
}
|
||||
|
||||
@@ -56,10 +76,16 @@ target "_labels" {
|
||||
"org.opencontainers.image.title" = "vLLM"
|
||||
"org.opencontainers.image.description" = "vLLM: A high-throughput and memory-efficient inference and serving engine for LLMs"
|
||||
"org.opencontainers.image.licenses" = "Apache-2.0"
|
||||
"org.opencontainers.image.revision" = COMMIT
|
||||
"org.opencontainers.image.revision" = VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")
|
||||
"org.opencontainers.image.version" = VLLM_IMAGE_TAG
|
||||
"org.opencontainers.image.url" = VLLM_BUILD_URL
|
||||
"ai.vllm.build.commit" = VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")
|
||||
"ai.vllm.build.pipeline" = VLLM_BUILD_PIPELINE
|
||||
"ai.vllm.build.url" = VLLM_BUILD_URL
|
||||
"ai.vllm.image.tag" = VLLM_IMAGE_TAG
|
||||
}
|
||||
annotations = [
|
||||
"index,manifest:org.opencontainers.image.revision=${COMMIT}",
|
||||
"index,manifest:org.opencontainers.image.revision=${VLLM_BUILD_COMMIT != "unknown" ? VLLM_BUILD_COMMIT : (COMMIT != "" ? COMMIT : "unknown")}",
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 156 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 182 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 193 KiB |
+2
-2
@@ -163,7 +163,7 @@ Running with a local file:
|
||||
|
||||
```bash
|
||||
vllm run-batch \
|
||||
-i offline_inference/openai_batch/openai_example_batch.jsonl \
|
||||
-i features/openai_batch/openai_example_batch.jsonl \
|
||||
-o results.jsonl \
|
||||
--model meta-llama/Meta-Llama-3-8B-Instruct
|
||||
```
|
||||
@@ -172,7 +172,7 @@ Using remote file:
|
||||
|
||||
```bash
|
||||
vllm run-batch \
|
||||
-i https://raw.githubusercontent.com/vllm-project/vllm/main/examples/offline_inference/openai_batch/openai_example_batch.jsonl \
|
||||
-i https://raw.githubusercontent.com/vllm-project/vllm/main/examples/features/openai_batch/openai_example_batch.jsonl \
|
||||
-o results.jsonl \
|
||||
--model meta-llama/Meta-Llama-3-8B-Instruct
|
||||
```
|
||||
|
||||
@@ -23,7 +23,7 @@ llm = LLM(model="ibm-granite/granite-3.1-8b-instruct", tensor_parallel_size=2)
|
||||
!!! note
|
||||
With tensor parallelism enabled, each process will read the whole model and split it into chunks, which makes the disk reading time even longer (proportional to the size of tensor parallelism).
|
||||
|
||||
You can convert the model checkpoint to a sharded checkpoint using [examples/offline_inference/save_sharded_state.py](../../examples/offline_inference/save_sharded_state.py). The conversion process might take some time, but later you can load the sharded checkpoint much faster. The model loading time should remain constant regardless of the size of tensor parallelism.
|
||||
You can convert the model checkpoint to a sharded checkpoint using [examples/features/sharded_state/load_sharded_state_offline.py](../../examples/features/sharded_state/load_sharded_state_offline.py). The conversion process might take some time, but later you can load the sharded checkpoint much faster. The model loading time should remain constant regardless of the size of tensor parallelism.
|
||||
|
||||
## Quantization
|
||||
|
||||
|
||||
@@ -42,7 +42,7 @@ Traces can be visualized using <https://ui.perfetto.dev/>.
|
||||
|
||||
#### Offline Inference
|
||||
|
||||
Refer to [examples/offline_inference/simple_profiling.py](../../examples/offline_inference/simple_profiling.py) for an example.
|
||||
Refer to [examples/features/profiling/simple_profiling_offline.py](../../examples/features/profiling/simple_profiling_offline.py) for an example.
|
||||
|
||||
#### OpenAI Server
|
||||
|
||||
|
||||
@@ -167,7 +167,7 @@ Priority is **1 = highest** (tried first).
|
||||
|
||||
| Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. |
|
||||
| ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ |
|
||||
| `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256, 512 | ❌ | ❌ | ❌ | All | N/A |
|
||||
| `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 |
|
||||
|
||||
@@ -11,7 +11,7 @@ Automatic Prefix Caching (APC in short) caches the KV cache of existing queries,
|
||||
|
||||
Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an example:
|
||||
|
||||
[examples/offline_inference/automatic_prefix_caching.py](../../examples/offline_inference/automatic_prefix_caching.py)
|
||||
[examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py](../../examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py)
|
||||
|
||||
## Example workloads
|
||||
|
||||
|
||||
@@ -6,12 +6,12 @@ This directory contains examples for extending the context length of models usin
|
||||
|
||||
## Offline Inference Example
|
||||
|
||||
The [`context_extension.py`](../../examples/offline_inference/context_extension) script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat example.
|
||||
The [`context_extension.py`](../../examples/features/context_extension/context_extension_offline.py) script demonstrates how to extend the context length of a Qwen model using the YARN method (rope_parameters) and run a simple chat example.
|
||||
|
||||
### Usage
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/context_extension.py
|
||||
python examples/features/context_extension/context_extension_offline.py
|
||||
```
|
||||
|
||||
## OpenAI Online Method
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
# IndexCache
|
||||
|
||||
IndexCache reduces redundant top-k computation in DeepSeek-V3.2 (DSA) models by caching and reusing top-k indices across layers.
|
||||
|
||||
## Background
|
||||
|
||||
DeepSeek-V3.2 uses a DeepSeek Sparse Attention (DSA) mechanism where top-k token selection is computed per layer. For deep models with many layers, this computation can be expensive. IndexCache allows skipping redundant top-k computations by reusing indices from previous layers.
|
||||
|
||||
See: [IndexCache Paper](https://arxiv.org/abs/2603.12201)
|
||||
|
||||
## Usage
|
||||
|
||||
### CLI
|
||||
|
||||
```bash
|
||||
vllm serve deepseek-ai/DeepSeek-V3.2 \
|
||||
--hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' ...
|
||||
```
|
||||
|
||||
### Configuration Reference
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|----------------------|------|---------|--------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| `use_index_cache` | bool | false | Enable IndexCache. Must be set to true to use this feature |
|
||||
| `index_topk_freq` | int | 1 | Frequency (in layers) at which top-k is computed. 1 = compute on every layer (disabled), 4 = compute on 1/4 of layers |
|
||||
| `index_topk_pattern` | str | null | Per-layer F/S pattern. Overrides index_topk_freq if set. Each character maps to one DSA layer: F = Full, S = Shared |
|
||||
|
||||
### Configuration Examples
|
||||
|
||||
**Using `index_topk_freq`** (compute every N layers):
|
||||
|
||||
```bash
|
||||
vllm serve deepseek-ai/DeepSeek-V3.2 \
|
||||
--hf-overrides '{"use_index_cache": true, "index_topk_freq": 4}' ...
|
||||
```
|
||||
|
||||
**Using `index_topk_pattern`** (explicit per-layer control):
|
||||
|
||||
```bash
|
||||
# custom pattern for 61 layers: F = compute, S = reuse
|
||||
vllm serve deepseek-ai/DeepSeek-V3.2 \
|
||||
--hf-overrides '{"use_index_cache": true, "index_topk_pattern": "FFSFSSSFSSFFFSSSFFFSFSSSSSSFFSFFSFFSSFFFFFFSFFFFFSFFSSSSSSFSF"}'
|
||||
```
|
||||
|
||||
## How It Works
|
||||
|
||||
1. When IndexCache is enabled, layers marked with `"F"` (Full) calculate and store top-k indices
|
||||
2. Subsequent layers marked with `"S"` (Shared) receive the cached indices from the previous layer instead of recomputing
|
||||
3. The cached indices are passed through the layer stack, reducing total computation
|
||||
|
||||
## Requirements
|
||||
|
||||
- DeepSeek-V3.2 or compatible DSA model
|
||||
- `use_index_cache: true` via `--hf-overrides`
|
||||
@@ -47,7 +47,7 @@ the third parameter is the path to the LoRA adapter.
|
||||
)
|
||||
```
|
||||
|
||||
Check out [examples/offline_inference/multilora_inference.py](../../examples/offline_inference/multilora_inference.py) for an example of how to use LoRA adapters with the async engine and how to use more advanced configuration options.
|
||||
Check out [examples/features/lora/multilora_offline.py](../../examples/features/lora/multilora_offline.py) for an example of how to use LoRA adapters with the async engine and how to use more advanced configuration options.
|
||||
|
||||
## Serving LoRA Adapters
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ To input multi-modal data, follow this schema in [vllm.inputs.EmbedsPrompt][]:
|
||||
|
||||
You can pass prompt embeddings from Hugging Face Transformers models to the `'prompt_embeds'` field of the prompt embedding dictionary, as shown in the following examples:
|
||||
|
||||
[examples/offline_inference/prompt_embed_inference.py](../../examples/offline_inference/prompt_embed_inference.py)
|
||||
[examples/features/prompt_embed/prompt_embed_offline.py](../../examples/features/prompt_embed/prompt_embed_offline.py)
|
||||
|
||||
## Online Serving
|
||||
|
||||
@@ -41,4 +41,4 @@ vllm serve meta-llama/Llama-3.2-1B-Instruct --runner generate \
|
||||
|
||||
Then, you can use the OpenAI client as follows:
|
||||
|
||||
[examples/online_serving/prompt_embed_inference_with_openai_client.py](../../examples/online_serving/prompt_embed_inference_with_openai_client.py)
|
||||
[examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py](../../examples/features/prompt_embed/prompt_embed_inference_with_openai_client.py)
|
||||
|
||||
@@ -13,6 +13,7 @@ vLLM currently supports the following reasoning models:
|
||||
|
||||
| Model Series | Parser Name | Structured Output Support | Tool Calling |
|
||||
| ------------ | ----------- | ---------------- | ----------- |
|
||||
| [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ |
|
||||
| [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ |
|
||||
| [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ |
|
||||
| [ERNIE-4.5-VL series](https://huggingface.co/baidu/ERNIE-4.5-VL-28B-A3B-PT) | `ernie45` | `json`, `regex` | ❌ |
|
||||
|
||||
@@ -32,7 +32,7 @@ depend on your model family, traffic pattern, hardware, and sampling settings.
|
||||
| Suffix decoding | Low to medium gain | Medium gain | No extra draft model; dynamic speculation depth. |
|
||||
|
||||
For reproducible measurements in your environment, use
|
||||
[`examples/offline_inference/spec_decode.py`](../../../examples/offline_inference/spec_decode.py)
|
||||
[`examples/features/speculative_decoding/spec_decode_offline.py`](../../../examples/features/speculative_decoding/spec_decode_offline.py)
|
||||
or the [benchmark CLI guide](../../benchmarking/cli.md).
|
||||
|
||||
## `--speculative-config` schema
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# EAGLE Draft Models
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/offline_inference/spec_decode.py](../../../examples/offline_inference/spec_decode.py)
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/features/speculative_decoding/spec_decode_offline.py](../../../examples/features/speculative_decoding/spec_decode_offline.py)
|
||||
|
||||
## Eagle Drafter Example
|
||||
|
||||
|
||||
@@ -165,7 +165,7 @@ As an example, we can use to define a specific format of simplified SQL queries:
|
||||
print(completion.choices[0].message.content)
|
||||
```
|
||||
|
||||
See also: [full example](../examples/online_serving/structured_outputs.md)
|
||||
See also: [full example](../../examples/features/structured_outputs/README.md)
|
||||
|
||||
## Reasoning Outputs
|
||||
|
||||
@@ -208,7 +208,7 @@ Note that you can use reasoning with any provided structured outputs feature. Th
|
||||
print("content: ", completion.choices[0].message.content)
|
||||
```
|
||||
|
||||
See also: [full example](../examples/online_serving/structured_outputs.md)
|
||||
See also: [full example](../../examples/features/structured_outputs/README.md)
|
||||
|
||||
!!! note
|
||||
When using Qwen3 Coder models with reasoning enabled, structured outputs might become disabled if the reasoning content does not get parsed into the `reasoning` field separately (v0.11.2+).
|
||||
@@ -304,7 +304,7 @@ Step #2: explanation="Next, let's isolate 'x' by dividing both sides of the equa
|
||||
Answer: x = -29/8
|
||||
```
|
||||
|
||||
An example of using `structural_tag` can be found here: [examples/online_serving/structured_outputs](../../examples/online_serving/structured_outputs)
|
||||
An example of using `structural_tag` can be found here: [examples/features/structured_outputs](../../examples/features/structured_outputs/README.md)
|
||||
|
||||
## Offline Inference
|
||||
|
||||
@@ -339,4 +339,4 @@ shown below:
|
||||
print(outputs[0].outputs[0].text)
|
||||
```
|
||||
|
||||
See also: [full example](../examples/online_serving/structured_outputs.md)
|
||||
See also: [full example](../../examples/features/structured_outputs/structured_outputs_offline.py)
|
||||
|
||||
@@ -369,6 +369,16 @@ Flags:
|
||||
* For non-reasoning: `--tool-call-parser hunyuan_a13b`
|
||||
* For reasoning: `--tool-call-parser hunyuan_a13b --reasoning-parser hunyuan_a13b`
|
||||
|
||||
### Cohere Command A Reasoning (`cohere_command3`)
|
||||
|
||||
Supported models:
|
||||
|
||||
* [`CohereLabs/command-a-reasoning-08-2025`](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025)
|
||||
|
||||
Flags: `--tool-call-parser cohere_command3 --reasoning-parser cohere_command3`
|
||||
|
||||
Note: the Cohere tool parser requires the `cohere_melody` package, which is not installed by default. Before using this parser please install the [cohere_melody](https://pypi.org/project/cohere-melody/) package.
|
||||
|
||||
### LongCat-Flash-Chat Models (`longcat`)
|
||||
|
||||
Supported models:
|
||||
|
||||
@@ -38,8 +38,20 @@ class MockCustomOp:
|
||||
return decorator
|
||||
|
||||
|
||||
class MockPluggableLayer:
|
||||
@staticmethod
|
||||
def register(name):
|
||||
def decorator(cls):
|
||||
return cls
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
mock_if_no_torch("vllm._C", MagicMock())
|
||||
mock_if_no_torch("vllm.model_executor.custom_op", MagicMock(CustomOp=MockCustomOp))
|
||||
mock_if_no_torch(
|
||||
"vllm.model_executor.custom_op",
|
||||
MagicMock(CustomOp=MockCustomOp, PluggableLayer=MockPluggableLayer),
|
||||
)
|
||||
mock_if_no_torch(
|
||||
"vllm.utils.torch_utils", MagicMock(direct_register_custom_op=lambda *a, **k: None)
|
||||
)
|
||||
|
||||
@@ -101,7 +101,7 @@ vllm serve /path/to/sharded/model \
|
||||
--model-loader-extra-config '{"pattern":"custom-model-rank-{rank}-part-{part}.safetensors"}'
|
||||
```
|
||||
|
||||
To create sharded model files, you can use the script provided in [examples/offline_inference/save_sharded_state.py](../../../examples/offline_inference/save_sharded_state.py). This script demonstrates how to save a model in the sharded format that is compatible with the Run:ai Model Streamer sharded loader.
|
||||
To create sharded model files, you can use the script provided in [examples/features/sharded_state/save_sharded_state_offline.py](../../../examples/features/sharded_state/save_sharded_state_offline.py). This script demonstrates how to save a model in the sharded format that is compatible with the Run:ai Model Streamer sharded loader.
|
||||
|
||||
The sharded loader supports all the same tunable parameters as the regular Run:ai Model Streamer, including `concurrency` and `memory_limit`. These can be configured in the same way:
|
||||
|
||||
|
||||
@@ -378,6 +378,7 @@ th {
|
||||
| `BloomForCausalLM` | BLOOM, BLOOMZ, BLOOMChat | `bigscience/bloom`, `bigscience/bloomz`, etc. | | ✅︎ |
|
||||
| `ChatGLMModel`, `ChatGLMForConditionalGeneration` | ChatGLM | `zai-org/chatglm2-6b`, `zai-org/chatglm3-6b`, `thu-coai/ShieldLM-6B-chatglm3`, etc. | ✅︎ | ✅︎ |
|
||||
| `CohereForCausalLM`, `Cohere2ForCausalLM` | Command-R, Command-A | `CohereLabs/c4ai-command-r-v01`, `CohereLabs/c4ai-command-r7b-12-2024`, `CohereLabs/c4ai-command-a-03-2025`, `CohereLabs/command-a-reasoning-08-2025`, etc. | ✅︎ | ✅︎ |
|
||||
| `CohereMoeForCausalLM` | Command (MoE) | (model checkpoints loaded with `trust_remote_code=True`) | ✅︎ | ✅︎ |
|
||||
| `CwmForCausalLM` | CWM | `facebook/cwm`, etc. | ✅︎ | ✅︎ |
|
||||
| `DbrxForCausalLM` | DBRX | `databricks/dbrx-base`, `databricks/dbrx-instruct`, etc. | | ✅︎ |
|
||||
| `DeciLMForCausalLM` | DeciLM | `nvidia/Llama-3_3-Nemotron-Super-49B-v1`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -439,7 +440,7 @@ th {
|
||||
| `Mamba2ForCausalLM` | Mamba2 | `mistralai/Mamba-Codestral-7B-v0.1`, etc. | | ✅︎ |
|
||||
| `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiMoV2FlashForCausalLM` | MiMoV2Flash | `XiaomiMiMo/MiMo-V2-Flash`, etc. | | ✅︎ |
|
||||
| `MiMoV2ProForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
|
||||
| `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ |
|
||||
| `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | |
|
||||
|
||||
@@ -16,7 +16,7 @@ For MoE models, when any requests are in progress in any rank, we must ensure th
|
||||
|
||||
In all cases, it is beneficial to load-balance requests between DP ranks. For online deployments, this balancing can be optimized by taking into account the state of each DP engine - in particular its currently scheduled and waiting (queued) requests, and KV cache state. Each DP engine has an independent KV cache, and the benefit of prefix caching can be maximized by directing prompts intelligently.
|
||||
|
||||
This document focuses on online deployments (with the API server). DP + EP is also supported for offline usage (via the LLM class), for an example see [examples/offline_inference/data_parallel.py](../../examples/offline_inference/data_parallel.py).
|
||||
This document focuses on online deployments (with the API server). DP + EP is also supported for offline usage (via the LLM class), for an example see [examples/features/data_parallel/data_parallel_offline.py](../../examples/features/data_parallel/data_parallel_offline.py).
|
||||
|
||||
There are two distinct modes supported for online deployments - self-contained with internal load balancing, or externally per-rank process deployment and load balancing.
|
||||
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
# What is Layerwise (Re)loading?
|
||||
|
||||
Layerwise reloading is the system used to handle the loading of new weight data into existing weight data destinations without triggering recompilation of the cuda graph and other runtime artifacts. This system is used to enable [QeRL](https://arxiv.org/pdf/2510.11696)-style post training flows, where full-precision trainer weights are quantized and loaded into a target vLLM instance for fast, high-exploration rollouts. The core implementation can be found in [layerwise.py](../../vllm/model_executor/model_loader/reload/layerwise.py).
|
||||
|
||||

|
||||
|
||||
## Layerwise Reloading for QeRL
|
||||
|
||||
In order to load new weights into existing weight data destinations, a weight must undergo the following operations:
|
||||
|
||||
- Transfer: weights must be transferred from trainer model to target node/device
|
||||
- Fuse: weight partitions must be fused, for example qkv/gate_up
|
||||
- Process: this typically means online quantization and kernel-specific padding or striding
|
||||
- Shard: weights must be sharded according to the selected parallelism strategy
|
||||
- Copy: weights must be copied into the existing weight data destinations
|
||||
|
||||
Layerwise reloading achieves this using the following steps:
|
||||
|
||||
1. Weights are **transferred** from the trainer to the target (see [weight_transfer](weight_transfer/README.md))
|
||||
2. Weights loaded via `model.load_weights`, during which they are **sharded** and **fused**
|
||||
3. Weights are **processed** in an online fashion as soon as all of a layer's weights are loaded
|
||||
4. Weights are **copied** into the existing weight data destinations
|
||||
|
||||
For more information on implementation, see [Low Level `layerwise` API](#low-level-layerwise-api).
|
||||
|
||||
## Layerwise Loading with Online Quantization
|
||||
|
||||
Online quantization refers to when a user provides full precision weights and those weights are quantized on-the-fly as they are loaded into the model. The layerwise reloading system handles this by treating online quantization as a **processing** step, which is then handled in an online way both during first-time load and during reload. A typical online quantization method implementation should look like this:
|
||||
|
||||
```python
|
||||
class Fp8OnlineLinearMethod(Fp8LinearMethod):
|
||||
"""Online version of Fp8LinearMethod which loads a full precision checkpoint
|
||||
and quantizes weights during loading."""
|
||||
|
||||
uses_meta_device: bool = True
|
||||
|
||||
def create_weights(self, layer: torch.nn.Module, ...):
|
||||
# weight is materialized and processed during loading
|
||||
layer.weight = ModelWeightParameter(
|
||||
data=torch.empty(..., device="meta"),
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
# set up online processing
|
||||
initialize_online_processing(layer)
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
if getattr(layer, "_already_called_process_weights_after_loading", False):
|
||||
return
|
||||
|
||||
layer.weight, layer.weight_scale = ops.scaled_fp8_quant(layer.weight)
|
||||
|
||||
# Prevent duplicate processing (e.g., during weight reload)
|
||||
layer._already_called_process_weights_after_loading = True
|
||||
```
|
||||
|
||||
## Example Usages
|
||||
|
||||
### High Level Weight Transfer API
|
||||
|
||||
The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Layerwise reloading is controlled by the `WeightTransferUpdateInfo.is_checkpoint_format` flag and is set to `True` by default.
|
||||
|
||||
### Mid Level `reload_weights` API
|
||||
|
||||
Layerwise reloading is also exposed via the `reload_weights` API. This interface can be called using the following code:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
llm = LLM("Qwen/Qwen3-0.6B")
|
||||
llm.collective_rpc("reload_weights")
|
||||
```
|
||||
|
||||
This interface also allows specifying a `weights_path` which can be used to select a checkpoint path to load from:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
# fine tuned model checkpoints for testing
|
||||
mul_path = "inference-optimization/Qwen3-0.6B-debug-multiply"
|
||||
add_path = "inference-optimization/Qwen3-0.6B-debug-add"
|
||||
|
||||
llm = LLM("Qwen/Qwen3-0.6B")
|
||||
llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_path})
|
||||
llm.generate("3 4 = ") # 12
|
||||
|
||||
llm.collective_rpc("reload_weights", kwargs={"weights_path": add_path})
|
||||
llm.generate("3 4 = ") # 7
|
||||
```
|
||||
|
||||
Finally, a `weights_iterator` can be provided directly. This iterator can be lazy or eagerly defined.
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
|
||||
weights_iterator = [("q_proj", ...), ("k_proj", ...), ...]
|
||||
|
||||
llm = LLM("Qwen/Qwen3-0.6B")
|
||||
llm.collective_rpc("reload_weights", kwargs={"weights_iterator": weights_iterator})
|
||||
```
|
||||
|
||||
### Low Level `layerwise` API
|
||||
|
||||
[layerwise.py](../../vllm/model_executor/model_loader/reload/layerwise.py) Implements the following functions to execute its lifecycle:
|
||||
|
||||
| Function | Purpose | Quantized Reload | Online Quantization |
|
||||
| - | - | - | - |
|
||||
| `record_metadata_for_reloading` | Record tensor metadata so that layers can be restored on the meta device | Called by `BaseModelLoader` | Called by `BaseModelLoader` |
|
||||
| `restore_layer_on_meta` | Restore layer to model format at start of reload | Called by `initialize_layerwise_reload` | Not called. Online quantized weights already start on meta device via `...OnlineLinearMethod.create_weights` |
|
||||
| `initialize_online_processing` | Wrap weight loaders with the `online_process_loader` wrapper, which buffers weights until all layer weights have been loaded | Called by `initialize_layerwise_reload` | Called by `...OnlineLinearMethod.create_weights` |
|
||||
| `_layerwise_process` | Process layer once all weights are loaded | Called by `online_process_loader` during loading | Called by `online_process_loader` during loading |
|
||||
| `_copy_and_restore_kernel_tensors` | Copy processed weights into original tensor locations to affect compiled cuda graphs, etc. | Called by `_layerwise_process` after `process_weights_after_loading` | Not called. There is no compiled cuda graph yet |
|
||||
| `finalize_layerwise_processing` | Catch any layers which did not load all weights (for example attention weights or weights with padding) | Called by `BaseModelLoader` | Called by `BaseModelLoader` |
|
||||
|
||||
You can plug into this lifecycle directly by calling the `initialize_layerwise_reload`, loading weights, then calling `finalize_layerwise_processing`:
|
||||
|
||||
```python
|
||||
from vllm import LLM
|
||||
from vllm.model_executor.model_loader.reload import initialize_layerwise_reload, finalize_layerwise_processing
|
||||
|
||||
llm = LLM("Qwen/Qwen3-0.6B")
|
||||
|
||||
# this model path requires `VLLM_ENABLE_V1_MULTIPROCESSING=0` and is not stable
|
||||
model = llm.llm_engine.engine_core.engine_core.model_executor.driver_worker.worker.get_model()
|
||||
|
||||
# layerwise reload
|
||||
initialize_layerwise_reload(model)
|
||||
model.load_weights(...)
|
||||
finalize_layerwise_processing(model, llm.model_config)
|
||||
```
|
||||
|
||||
## Troubleshooting Excessive Memory Usage
|
||||
|
||||
Layerwise reloading allows users to incrementally load and process weights as they are loaded into the model. This system relies on buffering layer weights on device until all weights of a layer have been loaded. However, without offloading, this approach necessarily causes excessive buffering if weights are loaded out of order.
|
||||
|
||||
For this reason, users must take care as to the order of weights when they are reloading into the model. Weight should be loaded "in order", meaning that each layer's weights are fully loaded before beginning to load the next layer's weights. "Out of order" loading can cause layer weights to stay buffered while other layer weights are loading, leading to excessive memory usage. In the example below, q_proj, k_proj, v_proj, and up_proj are all buffered at the same time, using more memory than if up_proj was loaded after q_proj, k_proj and v_proj.
|
||||
|
||||
| Correct Loading | Incorrect Loading |
|
||||
| - | - |
|
||||
|  |  |
|
||||
|
||||
Users will see a warning like the one below if weights are loaded out-of-order.
|
||||
|
||||
```console
|
||||
WARNING [layerwise.py:198] Allocating 28.5 MB of device memory to buffers to load ["QKVParallelLinear", "MergedColumnParallelLinear"] layers. This extra memory usage can be avoided by ordering weights by their parent layer when reloading.
|
||||
```
|
||||
@@ -7,7 +7,7 @@ reproducible results:
|
||||
or enable [batch invariance](../features/batch_invariance.md) to make the outputs insensitive to scheduling.
|
||||
- In online mode, you can only enable [batch invariance](../features/batch_invariance.md).
|
||||
|
||||
Example: [examples/offline_inference/reproducibility.py](../../examples/offline_inference/reproducibility.py)
|
||||
Example: [examples/features/batch_invariance/reproducibility_offline.py](../../examples/features/batch_invariance/reproducibility_offline.py)
|
||||
|
||||
!!! warning
|
||||
|
||||
|
||||
+26
-5
@@ -138,14 +138,22 @@ When `--api-key` is configured, the following `/v1` endpoints require Bearer tok
|
||||
|
||||
- `/v1/models` - List available models
|
||||
- `/v1/chat/completions` - Chat completions
|
||||
- `/v1/chat/completions/batch` - Batch chat completions
|
||||
- `/v1/chat/completions/render` - Render chat completion requests
|
||||
- `/v1/completions` - Text completions
|
||||
- `/v1/completions/render` - Render completion requests
|
||||
- `/v1/embeddings` - Generate embeddings
|
||||
- `/v1/audio/transcriptions` - Audio transcription
|
||||
- `/v1/audio/translations` - Audio translation
|
||||
- `/v1/messages` - Anthropic-compatible messages API
|
||||
- `/v1/responses` - Response management
|
||||
- `/v1/messages/count_tokens` - Count tokens for Anthropic messages
|
||||
- `/v1/responses` - Create a response
|
||||
- `/v1/responses/{response_id}` - Retrieve a response
|
||||
- `/v1/responses/{response_id}/cancel` - Cancel a response
|
||||
- `/v1/score` - Scoring API
|
||||
- `/v1/rerank` - Reranking API
|
||||
- `/v1/load_lora_adapter` - Load a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`)
|
||||
- `/v1/unload_lora_adapter` - Unload a LoRA adapter (can alter model behavior; only available when `--enable-lora` is set and `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True`)
|
||||
|
||||
### Unprotected Endpoints (No API Key Required)
|
||||
|
||||
@@ -155,16 +163,23 @@ The following endpoints **do not require authentication** even when `--api-key`
|
||||
|
||||
- `/invocations` - SageMaker-compatible endpoint (routes to the same inference functions as `/v1` endpoints)
|
||||
- `/inference/v1/generate` - Generate completions
|
||||
- `/generative_scoring` - Generative scoring API
|
||||
- `/pooling` - Pooling API
|
||||
- `/classify` - Classification API
|
||||
- `/score` - Scoring API (non-`/v1` variant)
|
||||
- `/rerank` - Reranking API (non-`/v1` variant)
|
||||
|
||||
**Operational control endpoints (always enabled):**
|
||||
**Operational control endpoints (only when `"generate"` task is supported):**
|
||||
|
||||
- `/pause` - Pause generation (causes denial of service)
|
||||
- `/resume` - Resume generation
|
||||
- `/is_paused` - Check if generation is paused
|
||||
- `/scale_elastic_ep` - Trigger scaling operations
|
||||
- `/is_scaling_elastic_ep` - Check if scaling is in progress
|
||||
- `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
|
||||
- `/update_weights` - Update model weights (can alter model behavior)
|
||||
- `/get_world_size` - Get distributed world size
|
||||
- `/abort_requests` - Abort in-flight requests (only when `--tokens-only` is also set)
|
||||
|
||||
**Utility endpoints:**
|
||||
|
||||
@@ -207,9 +222,9 @@ These endpoints are only available when profiling is enabled and should only be
|
||||
|
||||
An attacker who can reach the vLLM HTTP server can:
|
||||
|
||||
1. **Bypass authentication** by using non-`/v1` endpoints like `/invocations`, `/inference/v1/generate`, `/pooling`, `/classify`, `/score`, or `/rerank` to run arbitrary inference without credentials
|
||||
2. **Cause denial of service** by calling `/pause` or `/scale_elastic_ep` without a token
|
||||
3. **Access operational controls** to manipulate server state (e.g., pausing generation)
|
||||
1. **Bypass authentication** by using non-`/v1` endpoints like `/invocations`, `/inference/v1/generate`, `/generative_scoring`, `/pooling`, `/classify`, `/score`, or `/rerank` to run arbitrary inference without credentials
|
||||
2. **Cause denial of service** by calling `/pause`, `/scale_elastic_ep`, or `/abort_requests` without a token
|
||||
3. **Access operational controls** to manipulate server state (e.g., pausing generation, updating model weights via `/update_weights`)
|
||||
4. **If `--enable-tokenizer-info-endpoint` is set:** Access sensitive tokenizer configuration including chat templates, which may reveal prompt engineering strategies or other implementation details
|
||||
5. **If `VLLM_SERVER_DEV_MODE=1` is set:** Execute arbitrary RPC commands via `/collective_rpc`, reset caches, put the engine to sleep, and access detailed server configuration
|
||||
|
||||
@@ -288,6 +303,12 @@ To disable the Python code interpreter specifically, omit `code_interpreter` fro
|
||||
|
||||
**Consider a custom implementation**: The GPT-OSS Python tool is a reference implementation. For production deployments, consider implementing a custom code execution sandbox with stricter isolation guarantees. See the [GPT-OSS documentation](https://github.com/openai/gpt-oss?tab=readme-ov-file#python) for guidance.
|
||||
|
||||
## Dynamic LoRA Loading
|
||||
|
||||
vLLM supports dynamically loading and unloading LoRA adapters at runtime via the `/v1/load_lora_adapter` and `/v1/unload_lora_adapter` API endpoints. This functionality is **not enabled by default** — it requires both `--enable-lora` and the environment variable `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True` to be set.
|
||||
|
||||
**Warning:** Dynamic LoRA loading is not a secure operation and should not be enabled in deployments exposed to untrusted clients. If you must enable dynamic LoRA loading, restrict access to the `/v1/load_lora_adapter` and `/v1/unload_lora_adapter` endpoints to trusted administrators only, using a reverse proxy or network-level access controls. Do not expose these endpoints to end users. For details on configuring LoRA adapters, see the [LoRA Adapters documentation](../features/lora.md).
|
||||
|
||||
## Reporting Security Vulnerabilities
|
||||
|
||||
If you believe you have found a security vulnerability in vLLM, please report it following the project's security policy. For more information on how to report security issues and the project's security policy, please see the [vLLM Security Policy](https://github.com/vllm-project/vllm/blob/main/SECURITY.md).
|
||||
|
||||
+1
-1
@@ -15,7 +15,7 @@ compares the generation time for two queries that share the same prefix
|
||||
but ask different questions.
|
||||
|
||||
Run:
|
||||
python examples/offline_inference/automatic_prefix_caching.py
|
||||
python examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py
|
||||
"""
|
||||
|
||||
import time
|
||||
+1
-1
@@ -6,7 +6,7 @@ of a Qwen model using the YARN method (rope_parameters)
|
||||
and run a simple chat example.
|
||||
|
||||
Usage:
|
||||
python examples/offline_inference/context_extension.py
|
||||
python examples/features/context_extension/context_extension_offline.py
|
||||
"""
|
||||
|
||||
from vllm import LLM, RequestOutput, SamplingParams
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user