forked from Karylab-cklius/vllm
Compare commits
193
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
71475ffb12 | ||
|
|
17bcc8ee3c | ||
|
|
607f3a50a3 | ||
|
|
9167ef8dd7 | ||
|
|
d7045619c1 | ||
|
|
206aee7dd9 | ||
|
|
674e6ffdb6 | ||
|
|
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 | ||
|
|
2c06cf3486 | ||
|
|
e6f710a87f | ||
|
|
c245d35ff4 | ||
|
|
f8ac0c7cf0 | ||
|
|
ebf862c351 | ||
|
|
8d8062d0a7 | ||
|
|
985961345a | ||
|
|
706a04d34b | ||
|
|
22631f80a0 | ||
|
|
2cc008e7b4 | ||
|
|
5d5c776444 | ||
|
|
592ae6805c | ||
|
|
7b1bc0a3eb | ||
|
|
c0879d9483 | ||
|
|
f5f9878514 | ||
|
|
2ce95a761b | ||
|
+8 |
4d51588e23 | ||
|
|
32e45636e3 | ||
|
|
b39c266dae | ||
|
|
9558f43903 | ||
|
|
8cd174fa35 | ||
|
|
c798593f0d | ||
|
|
12a3f6454b | ||
|
|
60cd878a3b | ||
|
|
1e9f19ca3f | ||
|
|
6646c0c7e0 | ||
|
|
95995bbef8 | ||
|
|
07351e0883 | ||
|
|
428b988c98 | ||
|
|
e54894fc85 | ||
|
|
bc2ae5a3d6 | ||
|
|
a474da2813 | ||
|
|
ce6a199ecc | ||
|
|
f88763efc3 | ||
|
|
333529deae | ||
|
|
8825608205 | ||
|
|
095d2f87e8 | ||
|
|
21792520e7 | ||
|
|
5e11b40365 | ||
|
|
f768b4473e | ||
|
|
914d0464c1 | ||
|
|
9f771b3ab9 | ||
|
|
c9d3c6e6af | ||
|
|
51adca74e6 | ||
|
|
e8eb0490ce | ||
|
|
e8ee2a78db | ||
|
|
2ec18f5df4 | ||
|
|
6dec49f27e | ||
|
|
b5587e1013 | ||
|
|
9ad5abe772 | ||
|
|
7d3195ea9f | ||
|
|
512f522192 | ||
|
|
4c34b2f6fc | ||
|
|
cf8a613a87 | ||
|
|
01acf96c6f | ||
|
|
079a4cf399 | ||
|
|
9744b699ba | ||
|
|
c662b4359e | ||
|
|
100c7b65e7 | ||
|
|
56bdf85e10 | ||
|
|
eba73068ea | ||
|
|
e9f331d72e | ||
|
|
c9bf77df92 | ||
|
|
3041344287 | ||
|
|
92762edc53 | ||
|
|
626daa2076 | ||
|
|
fe85a92e86 | ||
|
|
62b1bbe470 | ||
|
|
fa4b70555b | ||
|
|
447c372ac5 | ||
|
|
ff2c2bd80a | ||
|
|
cde8d24710 | ||
|
|
4a6dd1c3cc | ||
|
|
7ff65b1900 | ||
|
|
7f95a66cbf | ||
|
|
1b1c01de39 | ||
|
|
e9ba519f45 | ||
|
|
5ef33ab250 | ||
|
|
1c2c1eb8b9 | ||
|
|
8824f50f1f | ||
|
|
0098db9ec1 | ||
|
|
53ecc807c0 | ||
|
|
b7a2605020 | ||
|
|
d0009ddb0b | ||
|
|
424033f4fc | ||
|
|
da1e7311ca | ||
|
|
01cb41dcf5 | ||
|
|
2f314bc5e6 | ||
|
|
2196bac135 | ||
|
|
4b7869d6bc | ||
|
|
4a79262e0f | ||
|
|
3ed5231c6a | ||
|
|
9c2492e501 | ||
|
|
fe57be7809 | ||
|
|
8317cedc77 | ||
|
|
98a242ff61 | ||
|
|
e4ee48da2d | ||
|
|
342c58bc54 | ||
|
|
fe9c3d6c5f | ||
|
|
ccaf5ffaa3 | ||
|
|
0283f303d8 | ||
|
|
ac58e2a170 | ||
|
|
b8401a9bf4 | ||
|
|
9c271f9403 | ||
|
|
22fa63cfe8 | ||
|
|
8f87eb4622 | ||
|
|
cfa49213d7 | ||
|
|
29f64c5f5e | ||
|
|
eb6661d522 | ||
|
|
d622e27d2b | ||
|
|
5f76b3fb30 | ||
|
|
809d83c2dc | ||
|
|
33ef1941e2 | ||
|
|
a4905133f3 | ||
|
|
ecbe42e991 | ||
|
|
a250f1bd5f | ||
|
|
04eac6ba24 | ||
|
|
9047288b68 | ||
|
|
ed6d30377d | ||
|
|
6aa057c9d7 | ||
|
|
a2bd09c960 | ||
|
|
123674879e | ||
|
|
4254aeb56f | ||
|
|
aad88f8486 | ||
|
|
0210024ae7 | ||
|
|
4eafc72928 | ||
|
|
6d09769700 | ||
|
|
4506319a28 | ||
|
|
9b60e2ffaa | ||
|
|
3951d3eacd | ||
|
|
6f2c71be8f | ||
|
|
2463f00fb6 | ||
|
|
f946659fff | ||
|
|
f90aa44662 | ||
|
|
cefa5281a7 | ||
|
|
46794958f0 | ||
|
|
6ff8dea075 | ||
|
|
583e6f2226 | ||
|
|
96a85c5750 | ||
|
|
9db4650e5e | ||
|
|
5e584ce9ec | ||
|
|
1842447c09 | ||
|
|
16688b26a6 | ||
|
|
6fbec8ed47 |
@@ -46,7 +46,7 @@ echo "Image not found, proceeding with build..."
|
||||
|
||||
# --- CUDA 13.0 for nightly builds ---
|
||||
# Nightly CI uses CUDA 13.0 while regular CI stays on CUDA 12.9
|
||||
NIGHTLY_CUDA_VERSION="13.0.0"
|
||||
NIGHTLY_CUDA_VERSION="13.0.2"
|
||||
NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
|
||||
NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ steps:
|
||||
'cd tests &&
|
||||
pytest -v -s lora/test_layers.py &&
|
||||
pytest -v -s lora/test_lora_checkpoints.py &&
|
||||
pytest -v -s lora/test_lora_functions.py &&
|
||||
(pytest -v -s lora/test_lora_functions.py --deselect="tests/lora/test_lora_functions.py::test_lora_functions_sync" --deselect="tests/lora/test_lora_functions.py::test_lora_functions_async" || true) &&
|
||||
pytest -v -s lora/test_lora_huggingface.py &&
|
||||
pytest -v -s lora/test_lora_manager.py &&
|
||||
pytest -v -s lora/test_lora_utils.py &&
|
||||
@@ -125,6 +125,5 @@ steps:
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'cd tests &&
|
||||
pytest -v -s lora/test_default_mm_loras.py &&
|
||||
pytest -v -s lora/test_qwen3_unembed.py &&
|
||||
pytest -v -s lora/test_qwenvl.py &&
|
||||
(pytest -v -s lora/test_qwen3_unembed.py || true) &&
|
||||
pytest -v -s lora/test_whisper.py'
|
||||
|
||||
@@ -1,3 +1,13 @@
|
||||
# CUDA architecture lists — following PyTorch RELEASE.md
|
||||
# (https://github.com/pytorch/pytorch/blob/main/RELEASE.md)
|
||||
# SM86 included for broader Ampere coverage; SM89 for marlin fp8 support
|
||||
env:
|
||||
CUDA_ARCH_X86: "7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
|
||||
# aarch64 only architectures: 8.7 for Orin, 11.0 for Thor (since CUDA 13)
|
||||
CUDA_ARCH_AARCH64: "8.0 8.7 8.9 9.0 10.0 11.0 12.0+PTX"
|
||||
CUDA_ARCH_X86_CU129: "7.5 8.0 8.6 8.9 9.0 10.0 12.0"
|
||||
CUDA_ARCH_AARCH64_CU129: "8.0 8.7 8.9 9.0 10.0 12.0"
|
||||
|
||||
steps:
|
||||
- input: "Provide Release version here"
|
||||
id: input-release-version
|
||||
@@ -14,12 +24,10 @@ steps:
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
- "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='8.7 8.9 9.0 10.0+PTX 12.0' --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=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"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
@@ -29,9 +37,7 @@ steps:
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "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"
|
||||
@@ -57,7 +63,7 @@ steps:
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=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"
|
||||
@@ -70,7 +76,7 @@ steps:
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "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"
|
||||
@@ -108,96 +114,95 @@ steps:
|
||||
depends_on: block-build-release-images
|
||||
allow_dependency_failure: true
|
||||
steps:
|
||||
- label: "Build release image - x86_64 - CUDA 12.9"
|
||||
- label: "Build release image - x86_64 - CUDA 13.0"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "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 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"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 12.9"
|
||||
- label: "Build release image - aarch64 - CUDA 13.0"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m) --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "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 push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 13.0"
|
||||
- label: "Build release image - x86_64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-cuda-13-0
|
||||
id: build-release-image-x86-cuda-12-9
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
|
||||
- "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 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)-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 13.0"
|
||||
- label: "Build release image - aarch64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-cuda-13-0
|
||||
id: build-release-image-arm64-cuda-12-9
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
# compute capability 12.0 for RTX-50 series / RTX PRO 6000 Blackwell, 12.1 for DGX Spark
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130"
|
||||
- "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 push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-ubuntu2404
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_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 push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-ubuntu2404
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.9.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0' --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "DOCKER_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 push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
|
||||
|
||||
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-x86-cuda-13-0-ubuntu2404
|
||||
id: build-release-image-x86-cuda-12-9-ubuntu2404
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
|
||||
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
- "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 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"
|
||||
|
||||
- label: "Build release image - aarch64 - CUDA 13.0 - Ubuntu 24.04"
|
||||
- label: "Build release image - aarch64 - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on: ~
|
||||
id: build-release-image-arm64-cuda-13-0-ubuntu2404
|
||||
id: build-release-image-arm64-cuda-12-9-ubuntu2404
|
||||
agents:
|
||||
queue: arm64_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.1 --build-arg UBUNTU_VERSION=24.04 --build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX 12.0 12.1' --build-arg INSTALL_KV_CONNECTORS=true --build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu24.04 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404 --target vllm-openai --progress plain -f docker/Dockerfile ."
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu130-ubuntu2404"
|
||||
- "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 push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
|
||||
|
||||
- block: "Build release image for x86_64 CPU"
|
||||
key: block-cpu-release-image-build
|
||||
@@ -238,7 +243,7 @@ steps:
|
||||
- group: "Publish release images"
|
||||
key: "publish-release-images"
|
||||
steps:
|
||||
- label: "Create multi-arch manifest - CUDA 12.9"
|
||||
- label: "Create multi-arch manifest - CUDA 13.0"
|
||||
depends_on:
|
||||
- build-release-image-x86
|
||||
- build-release-image-arm64
|
||||
@@ -250,7 +255,7 @@ steps:
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
|
||||
|
||||
- label: "Annotate release workflow - CUDA 12.9"
|
||||
- label: "Annotate release workflow - CUDA 13.0"
|
||||
depends_on:
|
||||
- create-multi-arch-manifest
|
||||
id: annotate-release-workflow
|
||||
@@ -259,19 +264,19 @@ steps:
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-release.sh"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 13.0"
|
||||
- label: "Create multi-arch manifest - CUDA 12.9"
|
||||
depends_on:
|
||||
- build-release-image-x86-cuda-13-0
|
||||
- build-release-image-arm64-cuda-13-0
|
||||
id: create-multi-arch-manifest-cuda-13-0
|
||||
- build-release-image-x86-cuda-12-9
|
||||
- build-release-image-arm64-cuda-12-9
|
||||
id: create-multi-arch-manifest-cuda-12-9
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
|
||||
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
|
||||
depends_on:
|
||||
- build-release-image-x86-ubuntu2404
|
||||
- build-release-image-arm64-ubuntu2404
|
||||
@@ -283,17 +288,17 @@ steps:
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-ubuntu2404"
|
||||
|
||||
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
|
||||
- label: "Create multi-arch manifest - CUDA 12.9 - Ubuntu 24.04"
|
||||
depends_on:
|
||||
- build-release-image-x86-cuda-13-0-ubuntu2404
|
||||
- build-release-image-arm64-cuda-13-0-ubuntu2404
|
||||
id: create-multi-arch-manifest-cuda-13-0-ubuntu2404
|
||||
- build-release-image-x86-cuda-12-9-ubuntu2404
|
||||
- build-release-image-arm64-cuda-12-9-ubuntu2404
|
||||
id: create-multi-arch-manifest-cuda-12-9-ubuntu2404
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu130-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu130-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu130-ubuntu2404"
|
||||
- "docker manifest create public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-x86_64-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-aarch64-cu129-ubuntu2404 --amend"
|
||||
- "docker manifest push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
|
||||
|
||||
- label: "Publish nightly multi-arch image to DockerHub"
|
||||
depends_on:
|
||||
@@ -313,16 +318,16 @@ steps:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
|
||||
- label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0"
|
||||
- label: "Publish nightly multi-arch image to DockerHub - CUDA 12.9"
|
||||
depends_on:
|
||||
- create-multi-arch-manifest-cuda-13-0
|
||||
- create-multi-arch-manifest-cuda-12-9
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/push-nightly-builds.sh cu130"
|
||||
- "bash .buildkite/scripts/push-nightly-builds.sh cu129"
|
||||
# Clean up old nightly builds (keep only last 14)
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-"
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh cu129-nightly-"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
|
||||
@@ -13,12 +13,12 @@ ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
|
||||
buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
|
||||
To download the wheel (by commit):
|
||||
\`\`\`
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_31_aarch64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
|
||||
(Optional) For CUDA 13.0:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu130-cp38-abi3-manylinux_2_35_aarch64.whl .
|
||||
(Optional) For CUDA 12.9:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_x86_64.whl .
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cu129-cp38-abi3-manylinux_2_31_aarch64.whl .
|
||||
|
||||
(Optional) For CPU:
|
||||
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
|
||||
@@ -33,8 +33,8 @@ To download and upload the image:
|
||||
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
|
||||
@@ -50,11 +50,11 @@ docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
docker push vllm/vllm-openai:latest-x86_64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130 vllm/vllm-openai:x86_64-cu130
|
||||
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:latest-x86_64-cu130
|
||||
docker tag vllm/vllm-openai:x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
|
||||
docker push vllm/vllm-openai:latest-x86_64-cu130
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu129 vllm/vllm-openai:x86_64-cu129
|
||||
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
|
||||
docker tag vllm/vllm-openai:x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
|
||||
docker push vllm/vllm-openai:latest-x86_64-cu129
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64 vllm/vllm-openai:aarch64
|
||||
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
|
||||
@@ -62,11 +62,11 @@ docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
docker push vllm/vllm-openai:latest-aarch64
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
|
||||
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130 vllm/vllm-openai:aarch64-cu130
|
||||
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker tag vllm/vllm-openai:aarch64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker push vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu129 vllm/vllm-openai:aarch64-cu129
|
||||
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker tag vllm/vllm-openai:aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
docker push vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
|
||||
## ROCm
|
||||
|
||||
@@ -104,11 +104,11 @@ docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${
|
||||
docker manifest push vllm/vllm-openai:latest
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
|
||||
|
||||
docker manifest rm vllm/vllm-openai:latest-cu130
|
||||
docker manifest create vllm/vllm-openai:latest-cu130 vllm/vllm-openai:latest-x86_64-cu130 vllm/vllm-openai:latest-aarch64-cu130
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu130 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu130
|
||||
docker manifest push vllm/vllm-openai:latest-cu130
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu130
|
||||
docker manifest rm vllm/vllm-openai:latest-cu129
|
||||
docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
|
||||
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
|
||||
docker manifest push vllm/vllm-openai:latest-cu129
|
||||
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
|
||||
|
||||
docker manifest rm vllm/vllm-openai-cpu:latest || true
|
||||
docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
|
||||
|
||||
@@ -29,7 +29,7 @@ if python3 -c "import torch; assert torch.version.hip" 2>/dev/null; then
|
||||
TORCH_INDEX_URL=""
|
||||
fi
|
||||
else
|
||||
TORCH_INDEX_URL="https://download.pytorch.org/whl/cu129"
|
||||
TORCH_INDEX_URL="https://download.pytorch.org/whl/cu130"
|
||||
fi
|
||||
echo ">>> Using PyTorch index: ${TORCH_INDEX_URL:-PyPI default}"
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ set -ex
|
||||
|
||||
BUCKET="vllm-wheels"
|
||||
INDICES_OUTPUT_DIR="indices"
|
||||
DEFAULT_VARIANT_ALIAS="cu129" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
|
||||
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/"
|
||||
|
||||
@@ -25,22 +25,100 @@ export PYTHONPATH=".."
|
||||
###############################################################################
|
||||
|
||||
cleanup_docker() {
|
||||
# Share the same lock with image pull to avoid cleanup/pull races on one node.
|
||||
local docker_lock="/tmp/docker-pull.lock"
|
||||
exec 9>"$docker_lock"
|
||||
flock 9
|
||||
|
||||
docker_root=$(docker info -f '{{.DockerRootDir}}')
|
||||
if [ -z "$docker_root" ]; then
|
||||
echo "Failed to determine Docker root directory." >&2
|
||||
exit 1
|
||||
flock -u 9
|
||||
return 1
|
||||
fi
|
||||
echo "Docker root directory: $docker_root"
|
||||
|
||||
disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
|
||||
threshold=70
|
||||
if [ "$disk_usage" -gt "$threshold" ]; then
|
||||
echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
|
||||
docker image prune -f
|
||||
docker volume prune -f && docker system prune --force --filter "until=72h" --all
|
||||
echo "Docker images and volumes cleanup completed."
|
||||
echo "Disk usage is above $threshold%. Running aggressive CI image cleanup..."
|
||||
cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 1
|
||||
else
|
||||
echo "Disk usage is below $threshold%. No cleanup needed."
|
||||
echo "Disk usage is below $threshold%. Checking old CI images anyway."
|
||||
cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 0
|
||||
fi
|
||||
echo "Old CI image cleanup completed."
|
||||
|
||||
flock -u 9
|
||||
}
|
||||
|
||||
cleanup_old_ci_images() {
|
||||
local repo_prefix="$1"
|
||||
local current_image_ref="$2"
|
||||
local ttl_hours="$3"
|
||||
local aggressive_cleanup="$4"
|
||||
|
||||
if [[ -z "$repo_prefix" || "$repo_prefix" == "/" ]]; then
|
||||
echo "Skip old-image cleanup: invalid repo prefix '${repo_prefix}'"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if ! [[ "$ttl_hours" =~ ^[0-9]+$ ]]; then
|
||||
echo "Invalid DOCKER_IMAGE_CLEANUP_HOURS='${ttl_hours}', fallback to 72"
|
||||
ttl_hours=72
|
||||
fi
|
||||
|
||||
local now_epoch cutoff_epoch
|
||||
now_epoch=$(date +%s)
|
||||
cutoff_epoch=$((now_epoch - ttl_hours * 3600))
|
||||
|
||||
local -a used_image_ids
|
||||
mapfile -t used_image_ids < <(docker ps -aq | xargs -r docker inspect --format '{{.Image}}' | sort -u)
|
||||
|
||||
local removed_count=0
|
||||
local examined_count=0
|
||||
declare -A seen_ids=()
|
||||
|
||||
while read -r image_ref image_id; do
|
||||
[[ -z "$image_ref" || -z "$image_id" ]] && continue
|
||||
((examined_count++))
|
||||
|
||||
# Keep the image this job is going to use.
|
||||
if [[ "$image_ref" == "$current_image_ref" ]]; then
|
||||
continue
|
||||
fi
|
||||
|
||||
# Avoid duplicate deletes when multiple tags point to same image id.
|
||||
if [[ -n "${seen_ids[$image_id]:-}" ]]; then
|
||||
continue
|
||||
fi
|
||||
seen_ids[$image_id]=1
|
||||
|
||||
# Never delete images that are used by any container on this node.
|
||||
if printf '%s\n' "${used_image_ids[@]}" | grep -qx "$image_id"; then
|
||||
continue
|
||||
fi
|
||||
|
||||
local created created_epoch
|
||||
created=$(docker image inspect -f '{{.Created}}' "$image_id" 2>/dev/null || true)
|
||||
[[ -z "$created" ]] && continue
|
||||
created_epoch=$(date -d "$created" +%s 2>/dev/null || true)
|
||||
[[ -z "$created_epoch" ]] && continue
|
||||
|
||||
if (( created_epoch < cutoff_epoch )) || [[ "$aggressive_cleanup" == "1" ]]; then
|
||||
if docker image rm -f "$image_id" >/dev/null 2>&1; then
|
||||
((removed_count++))
|
||||
fi
|
||||
fi
|
||||
done < <(docker image ls --no-trunc "$repo_prefix" --format '{{.Repository}}:{{.Tag}} {{.ID}}')
|
||||
|
||||
# Also trim old dangling layers; this is safe and does not remove referenced images.
|
||||
docker image prune -f --filter "until=${ttl_hours}h" >/dev/null 2>&1 || true
|
||||
|
||||
if [[ "$aggressive_cleanup" == "1" ]]; then
|
||||
echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} unused images under disk pressure."
|
||||
else
|
||||
echo "Examined ${examined_count} images under ${repo_prefix}, removed ${removed_count} old images (>${ttl_hours}h)."
|
||||
fi
|
||||
}
|
||||
|
||||
@@ -265,8 +343,6 @@ fi
|
||||
|
||||
remove_docker_container() {
|
||||
docker rm -f "${container_name}" || true
|
||||
docker image rm -f "${image_name}" || true
|
||||
docker system prune -f || true
|
||||
}
|
||||
trap remove_docker_container EXIT
|
||||
|
||||
|
||||
@@ -12,9 +12,7 @@ docker build -t "${image_name}" -f docker/Dockerfile.xpu .
|
||||
|
||||
# Setup cleanup
|
||||
remove_docker_container() {
|
||||
docker rm -f "${container_name}" || true;
|
||||
docker image rm -f "${image_name}" || true;
|
||||
docker system prune -f || true;
|
||||
docker rm -f "${container_name}" || true
|
||||
}
|
||||
trap remove_docker_container EXIT
|
||||
|
||||
|
||||
+31
-31
@@ -388,18 +388,18 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# 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 --------------------------------------------------------#
|
||||
|
||||
@@ -1647,18 +1647,18 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# 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 ----------------------------------------------------------#
|
||||
|
||||
@@ -1951,8 +1951,8 @@ steps:
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
#------------------------------------------------------- mi300 · quantization --------------------------------------------------------#
|
||||
|
||||
@@ -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:
|
||||
@@ -2930,18 +2930,18 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# Multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# 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 ----------------------------------------------------------#
|
||||
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
group: Disaggregated
|
||||
depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Distributed NixlConnector PD accuracy (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- 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)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- 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)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hybrid SSM NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
|
||||
|
||||
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
device: a100
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/v1/worker/kv_connector_model_runner_mixin.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
|
||||
@@ -88,9 +88,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,7 +106,7 @@ 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
|
||||
@@ -159,7 +158,7 @@ steps:
|
||||
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
|
||||
@@ -169,7 +168,7 @@ steps:
|
||||
# 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)
|
||||
device: a100
|
||||
@@ -194,7 +193,7 @@ 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
|
||||
@@ -222,94 +221,9 @@ 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"
|
||||
|
||||
- label: Distributed NixlConnector PD accuracy (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: DP EP Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- 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)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hyrbid SSM NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 20
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- HYBRID_SSM=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
|
||||
|
||||
- label: NixlConnector PD + Spec Decode acceptance (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
device: a100
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/v1/worker/kv_connector_model_runner_mixin.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
|
||||
- ./.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)
|
||||
timeout_in_minutes: 60
|
||||
|
||||
@@ -95,11 +95,13 @@ steps:
|
||||
- tests/kernels/moe/test_deepgemm.py
|
||||
- tests/kernels/moe/test_batched_deepgemm.py
|
||||
- tests/kernels/attention/test_deepgemm_attention.py
|
||||
- tests/quantization/test_cutlass_w4a16.py
|
||||
commands:
|
||||
- pytest -v -s kernels/quantization/test_block_fp8.py
|
||||
- pytest -v -s kernels/moe/test_deepgemm.py
|
||||
- pytest -v -s kernels/moe/test_batched_deepgemm.py
|
||||
- pytest -v -s kernels/attention/test_deepgemm_attention.py
|
||||
- pytest -v -s quantization/test_cutlass_w4a16.py
|
||||
|
||||
- label: Kernels (B200)
|
||||
timeout_in_minutes: 30
|
||||
|
||||
@@ -113,19 +113,19 @@ steps:
|
||||
- python3 basic/offline_inference/embed.py
|
||||
- python3 basic/offline_inference/score.py
|
||||
# for multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# 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)
|
||||
timeout_in_minutes: 20
|
||||
|
||||
@@ -31,8 +31,9 @@ steps:
|
||||
- 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:
|
||||
@@ -44,19 +45,19 @@ steps:
|
||||
#- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10 # TODO
|
||||
#- python3 basic/offline_inference/embed.py # TODO
|
||||
# for multi-modal models
|
||||
- python3 offline_inference/audio_language.py --seed 0
|
||||
- python3 offline_inference/vision_language.py --seed 0
|
||||
- python3 offline_inference/vision_language_multi_image.py --seed 0
|
||||
- python3 offline_inference/encoder_decoder_multimodal.py --model-type whisper --seed 0
|
||||
- python3 generate/multimodal/audio_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_offline.py --seed 0
|
||||
- python3 generate/multimodal/vision_language_multi_image_offline.py --seed 0
|
||||
- python3 generate/multimodal/encoder_decoder_multimodal_offline.py --model-type whisper --seed 0
|
||||
# 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)
|
||||
timeout_in_minutes: 45
|
||||
|
||||
@@ -69,9 +69,9 @@ steps:
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
- label: Transformers Backward Compatibility Models Test
|
||||
working_dir: "/vllm-workspace/"
|
||||
@@ -83,7 +83,7 @@ steps:
|
||||
- pytest -v -s tests/models/test_transformers.py
|
||||
- pytest -v -s tests/models/multimodal/processing/
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/offline_inference/basic/chat.py
|
||||
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
# Whisper needs spawn method to avoid deadlock
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
|
||||
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
|
||||
|
||||
@@ -28,6 +28,7 @@ steps:
|
||||
- pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
|
||||
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
|
||||
+4
-9
@@ -262,7 +262,7 @@ pull_request_rules:
|
||||
- files~=^docker/Dockerfile.xpu
|
||||
- files~=^\\.buildkite/intel_jobs/
|
||||
- files=\.buildkite/ci_config_intel.yaml
|
||||
- files=vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
|
||||
- files=vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
|
||||
- files=vllm/model_executor/kernels/linear/mixed_precision/xpu.py
|
||||
- files=vllm/model_executor/kernels/linear/mxfp8/xpu.py
|
||||
- files=vllm/model_executor/kernels/linear/scaled_mm/xpu.py
|
||||
@@ -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
|
||||
@@ -389,11 +388,7 @@ pull_request_rules:
|
||||
- files~=^tests/entrypoints/anthropic/.*tool.*
|
||||
- files~=^vllm/tool_parsers/
|
||||
- files=docs/features/tool_calling.md
|
||||
- files~=^examples/tool_chat_*
|
||||
- files=examples/offline_inference/chat_with_tools.py
|
||||
- files=examples/online_serving/openai_chat_completion_client_with_tools_required.py
|
||||
- files=examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py
|
||||
- files=examples/online_serving/openai_chat_completion_client_with_tools.py
|
||||
- files~=^examples/tool_calling/
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
|
||||
@@ -14,7 +14,7 @@ $python_executable -m pip install -r requirements/build/cuda.txt -r requirements
|
||||
# Limit the number of parallel jobs to avoid OOM
|
||||
export MAX_JOBS=1
|
||||
# Make sure release wheels are built for the following architectures
|
||||
export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
|
||||
export TORCH_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
|
||||
|
||||
bash tools/check_repo.sh
|
||||
|
||||
|
||||
+13
-7
@@ -34,10 +34,10 @@ install(CODE "set(CMAKE_INSTALL_LOCAL_ONLY TRUE)" ALL_COMPONENTS)
|
||||
# Supported python versions. These versions will be searched in order, the
|
||||
# first match will be selected. These should be kept in sync with setup.py.
|
||||
#
|
||||
set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13")
|
||||
set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13" "3.14")
|
||||
|
||||
# Supported AMD GPU architectures.
|
||||
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
|
||||
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
|
||||
|
||||
# ROCm installation prefix. Default to /opt/rocm but allow override via
|
||||
# -DROCM_PATH=/your/rocm/path when invoking cmake.
|
||||
@@ -94,12 +94,15 @@ find_package(Torch REQUIRED)
|
||||
# This check must happen after find_package(Torch) because that's when CMAKE_CUDA_COMPILER_VERSION gets defined
|
||||
if(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
|
||||
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
|
||||
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0;12.1")
|
||||
# starting from CUDA 12.9 and Blackwell (10.0), we use family-specific targets (10.0f, 12.0f, etc)
|
||||
# to support the whole generation without specifying all sub-architectures
|
||||
# see: https://developer.nvidia.com/blog/nvidia-blackwell-and-nvidia-cuda-12-9-introduce-family-specific-architecture-features/
|
||||
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0")
|
||||
elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
|
||||
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
|
||||
set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;12.0;12.1")
|
||||
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;10.1;10.3;12.0;12.1")
|
||||
else()
|
||||
set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0")
|
||||
set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
|
||||
endif()
|
||||
|
||||
#
|
||||
@@ -307,7 +310,9 @@ set(VLLM_EXT_SRC
|
||||
"csrc/torch_bindings.cpp")
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_EXT_SRC "csrc/minimax_reduce_rms_kernel.cu")
|
||||
list(APPEND VLLM_EXT_SRC
|
||||
"csrc/minimax_reduce_rms_kernel.cu"
|
||||
"csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
|
||||
|
||||
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
|
||||
|
||||
@@ -1048,7 +1053,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
list(APPEND VLLM_MOE_EXT_SRC
|
||||
"csrc/moe/moe_wna16.cu"
|
||||
"csrc/moe/grouped_topk_kernels.cu"
|
||||
"csrc/moe/router_gemm.cu")
|
||||
"csrc/moe/router_gemm.cu"
|
||||
"csrc/moe/topk_softplus_sqrt_kernels.cu")
|
||||
endif()
|
||||
|
||||
if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
|
||||
@@ -14,7 +14,7 @@ Easy, fast, and cheap LLM serving for everyone
|
||||
| <a href="https://docs.vllm.ai"><b>Documentation</b></a> | <a href="https://blog.vllm.ai/"><b>Blog</b></a> | <a href="https://arxiv.org/abs/2309.06180"><b>Paper</b></a> | <a href="https://x.com/vllm_project"><b>Twitter/X</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> |
|
||||
</p>
|
||||
|
||||
🔥 We have built a vllm website to help you get started with vllm. Please visit [vllm.ai](https://vllm.ai) to learn more.
|
||||
🔥 We have built a vLLM website to help you get started with vLLM. Please visit [vllm.ai](https://vllm.ai) to learn more.
|
||||
For events, please visit [vllm.ai/events](https://vllm.ai/events) to join us.
|
||||
|
||||
---
|
||||
@@ -50,7 +50,7 @@ vLLM is flexible and easy to use with:
|
||||
- Efficient multi-LoRA support for dense and MoE layers
|
||||
- Support for NVIDIA GPUs, AMD GPUs, and x86/ARM/PowerPC CPUs. Additionally, diverse hardware plugins such as Google TPUs, Intel Gaudi, IBM Spyre, Huawei Ascend, Rebellions NPU, Apple Silicon, MetaX GPU, and more.
|
||||
|
||||
vLLM seamlessly supports 200+ model architectures on HuggingFace, including:
|
||||
vLLM seamlessly supports 200+ model architectures on Hugging Face, including:
|
||||
|
||||
- Decoder-only LLMs (e.g., Llama, Qwen, Gemma)
|
||||
- Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS)
|
||||
|
||||
@@ -404,6 +404,7 @@ def _build_attention_metadata(
|
||||
query_start_loc=q_start_gpu,
|
||||
query_start_loc_cpu=q_start_cpu,
|
||||
seq_lens=seq_lens_gpu,
|
||||
seq_lens_cpu_upper_bound=seq_lens_cpu,
|
||||
_seq_lens_cpu=seq_lens_cpu,
|
||||
_num_computed_tokens_cpu=num_computed_tokens_cpu,
|
||||
slot_mapping=slot_mapping,
|
||||
|
||||
@@ -16,7 +16,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
|
||||
maybe_make_prepare_finalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
|
||||
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
@@ -22,7 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
|
||||
fp8_w8a8_moe_quant_config,
|
||||
nvfp4_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
|
||||
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
|
||||
CutlassExpertsFp4,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
|
||||
|
||||
@@ -13,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
|
||||
maybe_make_prepare_finalize,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
|
||||
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
|
||||
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import CutlassExpertsFp8
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import (
|
||||
fused_experts,
|
||||
fused_topk,
|
||||
|
||||
@@ -0,0 +1,324 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Benchmarks FP8 vs BF16 ViT attention via FlashInfer cuDNN backend.
|
||||
#
|
||||
# == Usage Examples ==
|
||||
#
|
||||
# Benchmark mode (default, FlashInfer CUDAGraph Bench)
|
||||
# python3 benchmark_vit_fp8_attn.py
|
||||
#
|
||||
# Profile mode (PyTorch profiler, saves TensorBoard traces):
|
||||
# python3 benchmark_vit_fp8_attn.py --profile
|
||||
# python3 benchmark_vit_fp8_attn.py --profile --profile-output-dir ./profile_traces
|
||||
#
|
||||
# Custom seq_lens:
|
||||
# python3 benchmark_vit_fp8_attn.py --seq-lens 4096 8192 16384
|
||||
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.profiler import ProfilerActivity, profile, record_function
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
# Qwen3-VL defaults
|
||||
NUM_HEADS = 16
|
||||
HEAD_DIM = 72
|
||||
DEFAULT_SEQ_LENS = [2304, 4096, 8192, 16384]
|
||||
|
||||
|
||||
def _setup_fp8_attention(num_heads: int, head_dim: int) -> tuple:
|
||||
"""Create FP8 and BF16 attention modules + workspace."""
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.multimodal import MultiModalConfig
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
_get_flashinfer_workspace_buffer,
|
||||
)
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
old_dtype = torch.get_default_dtype()
|
||||
torch.set_default_dtype(torch.bfloat16)
|
||||
|
||||
backend_patch = patch(
|
||||
"vllm.model_executor.layers.attention.mm_encoder_attention"
|
||||
".get_vit_attn_backend",
|
||||
return_value=AttentionBackendEnum.FLASHINFER,
|
||||
)
|
||||
|
||||
# FP8 attention
|
||||
mm_config_fp8 = MultiModalConfig(mm_encoder_attn_dtype="fp8")
|
||||
vllm_config_fp8 = VllmConfig()
|
||||
vllm_config_fp8.model_config = SimpleNamespace(multimodal_config=mm_config_fp8)
|
||||
with set_current_vllm_config(vllm_config_fp8), backend_patch:
|
||||
attn_fp8 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
# BF16 attention (no FP8)
|
||||
with set_current_vllm_config(VllmConfig()), backend_patch:
|
||||
attn_bf16 = MMEncoderAttention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_dim,
|
||||
prefix="visual.blocks.0.attn",
|
||||
).to("cuda")
|
||||
|
||||
torch.set_default_dtype(old_dtype)
|
||||
|
||||
workspace = _get_flashinfer_workspace_buffer()
|
||||
return attn_fp8, attn_bf16, workspace
|
||||
|
||||
|
||||
def _build_meta(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
fp8: bool,
|
||||
):
|
||||
"""Build cu_seqlens, max_seqlen, sequence_lengths."""
|
||||
from vllm.model_executor.layers.attention.mm_encoder_attention import (
|
||||
MMEncoderAttention,
|
||||
)
|
||||
from vllm.utils.math_utils import round_up
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
|
||||
cu_np = np.array([0, seq_len], dtype=np.int32)
|
||||
fp8_padded = num_heads * round_up(head_dim, 16) if fp8 else None
|
||||
|
||||
seq_lengths = MMEncoderAttention.maybe_compute_seq_lens(
|
||||
AttentionBackendEnum.FLASHINFER, cu_np, torch.device("cuda")
|
||||
)
|
||||
max_seqlen = torch.tensor(
|
||||
MMEncoderAttention.compute_max_seqlen(AttentionBackendEnum.FLASHINFER, cu_np),
|
||||
dtype=torch.int32,
|
||||
)
|
||||
cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
|
||||
AttentionBackendEnum.FLASHINFER,
|
||||
cu_np,
|
||||
num_heads * head_dim,
|
||||
1,
|
||||
torch.device("cuda"),
|
||||
fp8_padded_hidden_size=fp8_padded,
|
||||
)
|
||||
return cu_seqlens, max_seqlen, seq_lengths
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
seq_lens: list[int],
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
method: str,
|
||||
):
|
||||
"""Benchmark FP8 vs BF16 attention across seq_lens.
|
||||
|
||||
Uses FlashInfer GPU-level timing to measure pure kernel time,
|
||||
excluding CPU launch overhead.
|
||||
"""
|
||||
if method == "cupti":
|
||||
from flashinfer.testing import bench_gpu_time_with_cupti as bench_fn
|
||||
|
||||
bench_fn = partial(bench_fn, use_cuda_graph=True, cold_l2_cache=False)
|
||||
elif method == "cudagraph":
|
||||
from flashinfer.testing import (
|
||||
bench_gpu_time_with_cudagraph as bench_fn,
|
||||
)
|
||||
|
||||
bench_fn = partial(bench_fn, cold_l2_cache=False)
|
||||
else:
|
||||
raise ValueError(f"Invalid method: {method}")
|
||||
|
||||
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
|
||||
|
||||
print(f"Timing method: {method}")
|
||||
print(f"{'seq_len':>8} {'BF16 (us)':>12} {'FP8 (us)':>12} {'Speedup':>10}")
|
||||
print("-" * 46)
|
||||
|
||||
for seq_len in seq_lens:
|
||||
torch.manual_seed(42)
|
||||
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
|
||||
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
|
||||
cu_bf16 = cu_fp8.clone()
|
||||
|
||||
def bf16_fn(q=q, k=k, v=v, cu=cu_bf16, ms=max_s, sl=seq_l):
|
||||
attn_bf16._forward_flashinfer(q, k, v, cu, ms, sl)
|
||||
|
||||
def fp8_fn(q=q, k=k, v=v, cu=cu_fp8, ms=max_s, sl=seq_l):
|
||||
attn_fp8._forward_flashinfer(q, k, v, cu, ms, sl)
|
||||
|
||||
# bench_fn returns List[float] of per-iteration times in ms
|
||||
bf16_times = bench_fn(bf16_fn)
|
||||
fp8_times = bench_fn(fp8_fn)
|
||||
|
||||
bf16_us = np.median(bf16_times) * 1e3 # ms -> us
|
||||
fp8_us = np.median(fp8_times) * 1e3
|
||||
speedup = bf16_us / fp8_us if fp8_us > 0 else float("inf")
|
||||
|
||||
print(f"{seq_len:>8} {bf16_us:>12.1f} {fp8_us:>12.1f} {speedup:>9.2f}x")
|
||||
|
||||
|
||||
def _make_trace_handler(output_dir: str, worker_name: str, label: str):
|
||||
"""Create a trace handler that saves to TensorBoard and prints summary."""
|
||||
|
||||
def handler(prof):
|
||||
torch.profiler.tensorboard_trace_handler(output_dir, worker_name)(prof)
|
||||
print(f"\n{'=' * 80}")
|
||||
print(label)
|
||||
print(f"{'=' * 80}")
|
||||
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
def run_profile(
|
||||
seq_len: int,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
warmup: int,
|
||||
output_dir: str,
|
||||
):
|
||||
"""Profile FP8 vs BF16 attention with PyTorch profiler."""
|
||||
attn_fp8, attn_bf16, workspace = _setup_fp8_attention(num_heads, head_dim)
|
||||
|
||||
torch.manual_seed(42)
|
||||
q = torch.randn(
|
||||
seq_len,
|
||||
num_heads,
|
||||
head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
k = torch.randn_like(q)
|
||||
v = torch.randn_like(q)
|
||||
|
||||
cu_fp8, max_s, seq_l = _build_meta(seq_len, num_heads, head_dim, fp8=True)
|
||||
# we can reuse cu_fp8 for cu_bf16 since q, k, and v are contiguous
|
||||
cu_bf16 = cu_fp8.clone()
|
||||
|
||||
sched = torch.profiler.schedule(wait=0, warmup=warmup, active=1)
|
||||
|
||||
# Profile BF16 (warmup handled by profiler schedule)
|
||||
with profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
||||
schedule=sched,
|
||||
on_trace_ready=_make_trace_handler(
|
||||
output_dir,
|
||||
f"bf16_h{head_dim}_s{seq_len}",
|
||||
f"BF16 Attention (seq_len={seq_len}, heads={num_heads}, "
|
||||
f"head_dim={head_dim})",
|
||||
),
|
||||
) as prof_bf16:
|
||||
for _ in range(warmup + 1):
|
||||
with record_function("bf16_attention"):
|
||||
attn_bf16._forward_flashinfer(
|
||||
q.clone(), k.clone(), v.clone(), cu_bf16, max_s, seq_l
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
prof_bf16.step()
|
||||
|
||||
# Profile FP8 (warmup handled by profiler schedule)
|
||||
with profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
||||
schedule=sched,
|
||||
on_trace_ready=_make_trace_handler(
|
||||
output_dir,
|
||||
f"fp8_h{head_dim}_s{seq_len}",
|
||||
f"FP8 Attention (seq_len={seq_len}, heads={num_heads}, "
|
||||
f"head_dim={head_dim})",
|
||||
),
|
||||
) as prof_fp8:
|
||||
for _ in range(warmup + 1):
|
||||
with record_function("fp8_attention"):
|
||||
attn_fp8._forward_flashinfer(
|
||||
q.clone(), k.clone(), v.clone(), cu_fp8, max_s, seq_l
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
prof_fp8.step()
|
||||
|
||||
print(f"\nTensorBoard traces saved to: {output_dir}")
|
||||
print(f"View with: tensorboard --logdir={output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = FlexibleArgumentParser(description="Benchmark FP8 vs BF16 ViT attention.")
|
||||
parser.add_argument(
|
||||
"--seq-lens",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=DEFAULT_SEQ_LENS,
|
||||
help="Sequence lengths to benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-heads",
|
||||
type=int,
|
||||
default=NUM_HEADS,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--head-dim",
|
||||
type=int,
|
||||
default=HEAD_DIM,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--method",
|
||||
choices=["cupti", "cudagraph"],
|
||||
default="cudagraph",
|
||||
help="GPU timing method: cupti (CUPTI kernel timing) or "
|
||||
"cudagraph (CUDA graph capture/replay). Default: cudagraph",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Warmup iterations (profile mode only)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile",
|
||||
action="store_true",
|
||||
help="Run PyTorch profiler instead of benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-seq-len",
|
||||
type=int,
|
||||
default=8192,
|
||||
help="Sequence length for profiling (default: 8192)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profile-output-dir",
|
||||
type=str,
|
||||
default="./profile_traces",
|
||||
help="Output directory for TensorBoard traces (default: ./profile_traces)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.profile:
|
||||
run_profile(
|
||||
args.profile_seq_len,
|
||||
args.num_heads,
|
||||
args.head_dim,
|
||||
args.warmup,
|
||||
args.profile_output_dir,
|
||||
)
|
||||
else:
|
||||
run_benchmark(
|
||||
args.seq_lens,
|
||||
args.num_heads,
|
||||
args.head_dim,
|
||||
args.method,
|
||||
)
|
||||
@@ -20,7 +20,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
deepgemm
|
||||
GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
|
||||
GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03
|
||||
GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
|
||||
GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
@@ -120,6 +120,11 @@ if(DEEPGEMM_ARCHS)
|
||||
COMPONENT _deep_gemm_C
|
||||
FILES_MATCHING PATTERN "*.py")
|
||||
|
||||
install(DIRECTORY "${deepgemm_SOURCE_DIR}/deep_gemm/mega/"
|
||||
DESTINATION vllm/third_party/deep_gemm/mega
|
||||
COMPONENT _deep_gemm_C
|
||||
FILES_MATCHING PATTERN "*.py")
|
||||
|
||||
# Generate envs.py (normally generated by DeepGEMM's setup.py build step)
|
||||
file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/deep_gemm_envs.py"
|
||||
"# Pre-installed environment variables\npersistent_envs = dict()\n")
|
||||
|
||||
@@ -19,7 +19,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
flashmla
|
||||
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
|
||||
GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
|
||||
GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
|
||||
+82
-25
@@ -11,29 +11,74 @@
|
||||
namespace vllm {
|
||||
|
||||
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
|
||||
bool act_first>
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
|
||||
const scalar_t& y) {
|
||||
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
|
||||
const scalar_t& y,
|
||||
const float limit) {
|
||||
if constexpr (act_first) {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
gate = (scalar_t)fminf((float)gate, limit);
|
||||
up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
|
||||
}
|
||||
return ACT_FN(gate) * up;
|
||||
} else {
|
||||
scalar_t gate = x;
|
||||
scalar_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
|
||||
up = (scalar_t)fminf((float)up, limit);
|
||||
}
|
||||
return gate * ACT_FN(up);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
|
||||
bool act_first>
|
||||
bool act_first, bool HAS_CLAMP>
|
||||
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
|
||||
const packed_t& y) {
|
||||
return act_first ? packed_mul(PACKED_ACT_FN(x), y)
|
||||
: packed_mul(x, PACKED_ACT_FN(y));
|
||||
const packed_t& y,
|
||||
const float limit) {
|
||||
if constexpr (act_first) {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fminf(g.x, limit);
|
||||
g.y = fminf(g.y, limit);
|
||||
u.x = fmaxf(fminf(u.x, limit), -limit);
|
||||
u.y = fmaxf(fminf(u.y, limit), -limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(PACKED_ACT_FN(gate), up);
|
||||
} else {
|
||||
packed_t gate = x;
|
||||
packed_t up = y;
|
||||
if constexpr (HAS_CLAMP) {
|
||||
float2 g = cast_to_float2(gate);
|
||||
float2 u = cast_to_float2(up);
|
||||
g.x = fmaxf(fminf(g.x, limit), -limit);
|
||||
g.y = fmaxf(fminf(g.y, limit), -limit);
|
||||
u.x = fminf(u.x, limit);
|
||||
u.y = fminf(u.y, limit);
|
||||
gate = cast_to_packed<packed_t>(g);
|
||||
up = cast_to_packed<packed_t>(u);
|
||||
}
|
||||
return packed_mul(gate, PACKED_ACT_FN(up));
|
||||
}
|
||||
}
|
||||
|
||||
// Activation and gating kernel template.
|
||||
template <typename scalar_t, typename packed_t,
|
||||
scalar_t (*ACT_FN)(const scalar_t&),
|
||||
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
|
||||
bool use_vec, bool use_256b = false>
|
||||
bool use_vec, bool HAS_CLAMP, bool use_256b = false>
|
||||
__global__ void act_and_mul_kernel(
|
||||
scalar_t* __restrict__ out, // [..., d]
|
||||
const scalar_t* __restrict__ input, // [..., 2, d]
|
||||
const int d) {
|
||||
const int d, const float limit) {
|
||||
const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
|
||||
const scalar_t* y_ptr = x_ptr + d;
|
||||
scalar_t* out_ptr = out + blockIdx.x * d;
|
||||
@@ -58,8 +103,9 @@ __global__ void act_and_mul_kernel(
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
|
||||
x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
|
||||
x.elts[j], y.elts[j]);
|
||||
x.elts[j] =
|
||||
packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
|
||||
x.elts[j], y.elts[j], limit);
|
||||
}
|
||||
if constexpr (use_256b) {
|
||||
st256(x, &out_vec[i]);
|
||||
@@ -72,7 +118,8 @@ __global__ void act_and_mul_kernel(
|
||||
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
|
||||
const scalar_t x = VLLM_LDG(&x_ptr[idx]);
|
||||
const scalar_t y = VLLM_LDG(&y_ptr[idx]);
|
||||
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
|
||||
out_ptr[idx] =
|
||||
compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -151,8 +198,11 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
|
||||
// Launch activation and gating kernel.
|
||||
// Use ACT_FIRST (bool) indicating whether to apply the activation function
|
||||
// first.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
|
||||
// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
|
||||
// clamped (max only) and up input is clamped (both sides) before the
|
||||
// activation function is applied.
|
||||
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
|
||||
HAS_CLAMP, LIMIT) \
|
||||
auto dtype = input.scalar_type(); \
|
||||
int d = input.size(-1) / 2; \
|
||||
int64_t num_tokens = input.numel() / input.size(-1); \
|
||||
@@ -177,8 +227,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
} else { \
|
||||
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
|
||||
@@ -186,8 +236,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
} \
|
||||
} else { \
|
||||
@@ -197,8 +247,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
|
||||
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
|
||||
KERNEL<scalar_t>, \
|
||||
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
|
||||
ACT_FIRST, false><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
|
||||
ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
|
||||
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
|
||||
}); \
|
||||
}
|
||||
|
||||
@@ -206,7 +256,14 @@ void silu_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true);
|
||||
true, false, 0.0f);
|
||||
}
|
||||
|
||||
void silu_and_mul_clamp(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input, // [..., 2 * d]
|
||||
double limit) {
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
true, true, (float)limit);
|
||||
}
|
||||
|
||||
void mul_and_silu(torch::Tensor& out, // [..., d]
|
||||
@@ -215,21 +272,21 @@ void mul_and_silu(torch::Tensor& out, // [..., d]
|
||||
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
|
||||
// applies the silu to the latter half of the input.
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
|
||||
false);
|
||||
false, false, 0.0f);
|
||||
}
|
||||
|
||||
void gelu_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
|
||||
true);
|
||||
true, false, 0.0f);
|
||||
}
|
||||
|
||||
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
|
||||
torch::Tensor& input) // [..., 2 * d]
|
||||
{
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
|
||||
vllm::packed_gelu_tanh_kernel, true);
|
||||
LAUNCH_ACTIVATION_GATE_KERNEL(
|
||||
vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
|
||||
}
|
||||
|
||||
namespace vllm {
|
||||
|
||||
+18
-8
@@ -599,6 +599,11 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
|
||||
const int head_idx = (blockIdx.y * blockDim.x + threadIdx.x) * VEC_SIZE;
|
||||
// Find batch index within a block
|
||||
__shared__ int batch_idx[BLOCK_Y_SIZE];
|
||||
if (threadIdx.x == 0) {
|
||||
batch_idx[threadIdx.y] = -1;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int iter = 0; iter < cuda_utils::ceil_div(batch_size, int(blockDim.x));
|
||||
iter++) {
|
||||
int tid = iter * blockDim.x + threadIdx.x;
|
||||
@@ -611,16 +616,18 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
|
||||
}
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
__syncwarp();
|
||||
#endif
|
||||
__syncthreads();
|
||||
|
||||
if (head_idx >= head_dim || token_idx >= num_tokens) {
|
||||
// num_tokens may be an allocation upper bound when Python avoids a D2H sync.
|
||||
// Only tokens covered by the exact device-side cu_seq_lens are valid to
|
||||
// gather.
|
||||
const int batch = batch_idx[threadIdx.y];
|
||||
if (head_idx >= head_dim || token_idx >= num_tokens || batch < 0) {
|
||||
return;
|
||||
}
|
||||
const int inbatch_seq_idx = token_idx - cu_seq_lens[batch_idx[threadIdx.y]];
|
||||
const int block_idx = block_table[batch_idx[threadIdx.y] * num_blocks +
|
||||
inbatch_seq_idx / cache_block_size];
|
||||
const int inbatch_seq_idx = token_idx - cu_seq_lens[batch];
|
||||
const int block_idx =
|
||||
block_table[batch * num_blocks + inbatch_seq_idx / cache_block_size];
|
||||
const int64_t src_block_offset = block_idx * block_stride;
|
||||
const int64_t cache_inblock_offset =
|
||||
(inbatch_seq_idx % cache_block_size) * head_dim + head_idx;
|
||||
@@ -1490,6 +1497,9 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
|
||||
TORCH_CHECK(ql_nope.stride(2) == 1, "ql_nope must have stride 1 in dim 2");
|
||||
TORCH_CHECK(q_pe.stride(2) == 1, "q_pe must have stride 1 in dim 2");
|
||||
TORCH_CHECK(q_out.stride(2) == 1, "q_out must have stride 1 in dim 2");
|
||||
TORCH_CHECK(ql_nope.scalar_type() == at::ScalarType::Half ||
|
||||
ql_nope.scalar_type() == at::ScalarType::BFloat16,
|
||||
"ql_nope must be float16 or bfloat16 dtype");
|
||||
|
||||
if (num_tokens == 0) return;
|
||||
|
||||
@@ -1501,7 +1511,7 @@ void concat_mla_q(torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(ql_nope));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
|
||||
VLLM_DISPATCH_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
|
||||
vllm::ConcatMLAQKernel<scalar_t, 512><<<grid_size, block_size, 0, stream>>>(
|
||||
q_out.data_ptr<scalar_t>(), ql_nope.data_ptr<scalar_t>(),
|
||||
q_pe.data_ptr<scalar_t>(), num_tokens, num_heads, q_out.stride(0),
|
||||
|
||||
@@ -61,8 +61,23 @@
|
||||
#endif
|
||||
|
||||
#ifdef __aarch64__
|
||||
// Implementation copied from Arm Optimized Routines (expf AdvSIMD)
|
||||
// Implementation of neon_expf copied from Arm Optimized Routines (expf
|
||||
// AdvSIMD)
|
||||
// https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/expf.c
|
||||
//
|
||||
// Additional fast exponential intended for cases where outputs will be
|
||||
// downcasted to FP16 / BF16 (e.g. attention softmax). Accurate within 1 ULP
|
||||
// for FP16 Accurate within 1 ULP for BF16 for inputs in [-87.683, 88.376] &
|
||||
// clamps inputs outside this range to 0 / inf. Implementation is similar to
|
||||
// exp_u20, but:
|
||||
// - uses a third degree polynomial approximation for exp(r) instead of a
|
||||
// fifth degree one, with coefficients re-tuned.
|
||||
// - does not split natural log (ln) into high / low parts
|
||||
// - clamps exp(x) to 0 for x < -87.683113f and inf for x > 88.3762589f
|
||||
// exp(x) = 2^n (exp(r))
|
||||
// r = x - n*ln2, with n = round(x/ln2)
|
||||
// exp(r) ~ poly(r) = 1 + r + r^2 * (c3 + c2 * r)
|
||||
// n = round(x / ln2), r = x - n*ln2
|
||||
#include <limits>
|
||||
#define DEFINE_FAST_EXP \
|
||||
const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f); \
|
||||
@@ -106,7 +121,38 @@
|
||||
result.val[2] = neon_expf(vec.reg.val[2]); \
|
||||
result.val[3] = neon_expf(vec.reg.val[3]); \
|
||||
return vec_op::FP32Vec16(result); \
|
||||
};
|
||||
}; \
|
||||
const float32x4_t lower_bound = vdupq_n_f32(-0x1.5ebb82p+6f); \
|
||||
const float32x4_t upper_bound = vdupq_n_f32(0x1.61814ap+6f); \
|
||||
constexpr float ln2 = 0x1.62e43p-1f; \
|
||||
constexpr float f_c2 = 0x1.5592ecp-3f; \
|
||||
const float32x4_t f_c3 = vdupq_n_f32(0x1.017d34p-1f); \
|
||||
auto neon_expf_f16 = [&](float32x4_t values) __attribute__(( \
|
||||
always_inline)) { \
|
||||
const uint32x4_t lt_lower = vcltq_f32(values, lower_bound); \
|
||||
const uint32x4_t gt_upper = vcgtq_f32(values, upper_bound); \
|
||||
float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2)); \
|
||||
float32x4_t r = vfmsq_n_f32(values, n, ln2); \
|
||||
uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \
|
||||
float32x4_t r2 = vmulq_f32(r, r); \
|
||||
float32x4_t q = vfmaq_n_f32(f_c3, r, f_c2); \
|
||||
float32x4_t s = vaddq_f32(vdupq_n_f32(1.0f), r); \
|
||||
float32x4_t p = vfmaq_f32(s, q, r2); \
|
||||
float32x4_t y = \
|
||||
vreinterpretq_f32_u32(vaddq_u32(vreinterpretq_u32_f32(p), e)); \
|
||||
y = vbslq_f32(lt_lower, vdupq_n_f32(0.0f), y); \
|
||||
y = vbslq_f32(gt_upper, vdupq_n_f32(INFINITY), y); \
|
||||
return y; \
|
||||
}; \
|
||||
auto fast_exp_f16 = [&](const vec_op::FP32Vec16& vec) \
|
||||
__attribute__((always_inline)) { \
|
||||
float32x4x4_t result; \
|
||||
result.val[0] = neon_expf_f16(vec.reg.val[0]); \
|
||||
result.val[1] = neon_expf_f16(vec.reg.val[1]); \
|
||||
result.val[2] = neon_expf_f16(vec.reg.val[2]); \
|
||||
result.val[3] = neon_expf_f16(vec.reg.val[3]); \
|
||||
return vec_op::FP32Vec16(result); \
|
||||
};
|
||||
|
||||
#endif // __aarch64__
|
||||
|
||||
|
||||
@@ -1152,7 +1152,11 @@ class AttentionMainLoop {
|
||||
bool use_sink) {
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
DEFINE_FAST_EXP
|
||||
bool constexpr IsReducedPrecision =
|
||||
std::is_same_v<query_t, c10::BFloat16> ||
|
||||
std::is_same_v<query_t, c10::Half>;
|
||||
#endif
|
||||
|
||||
using prob_buffer_vec_t = typename VecTypeTrait<prob_buffer_t>::vec_t;
|
||||
static_assert(sizeof(prob_buffer_t) <= sizeof(logits_buffer_t));
|
||||
|
||||
@@ -1201,8 +1205,17 @@ class AttentionMainLoop {
|
||||
vec = vec - max_vec;
|
||||
|
||||
// compute exp
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
vec = fast_exp(vec);
|
||||
|
||||
#if defined(DEFINE_FAST_EXP)
|
||||
#ifdef __aarch64__
|
||||
if constexpr (IsReducedPrecision) {
|
||||
vec = fast_exp_f16(vec);
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
vec = fast_exp(vec);
|
||||
}
|
||||
|
||||
prob_buffer_vec_t output_vec(vec);
|
||||
output_vec.save(curr_prob_buffer_iter);
|
||||
#else
|
||||
@@ -1258,7 +1271,11 @@ class AttentionMainLoop {
|
||||
int32_t kv_tile_token_num, float softcap_scale) {
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
DEFINE_FAST_EXP
|
||||
bool constexpr IsReducedPrecision =
|
||||
std::is_same_v<query_t, c10::BFloat16> ||
|
||||
std::is_same_v<query_t, c10::Half>;
|
||||
#endif
|
||||
|
||||
float inv_softcap_scale = 1.0 / softcap_scale;
|
||||
vec_op::FP32Vec16 softcap_scale_vec(softcap_scale);
|
||||
vec_op::FP32Vec16 inv_softcap_scale_vec(inv_softcap_scale);
|
||||
@@ -1272,8 +1289,15 @@ class AttentionMainLoop {
|
||||
vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
|
||||
vec = vec * inv_softcap_scale_vec;
|
||||
|
||||
#ifdef DEFINE_FAST_EXP
|
||||
vec = fast_exp(vec);
|
||||
#if defined(DEFINE_FAST_EXP)
|
||||
#ifdef __aarch64__
|
||||
if constexpr (IsReducedPrecision) {
|
||||
vec = fast_exp_f16(vec);
|
||||
} else
|
||||
#endif
|
||||
{
|
||||
vec = fast_exp(vec);
|
||||
}
|
||||
vec_op::FP32Vec16 inv_vec = ones_vec / vec;
|
||||
vec = (vec - inv_vec) / (vec + inv_vec);
|
||||
#else
|
||||
|
||||
@@ -15,16 +15,12 @@
|
||||
#include <torch/all.h>
|
||||
namespace vec_op {
|
||||
|
||||
#ifdef RISCV_BF16_SUPPORT
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
|
||||
#else
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
|
||||
#endif
|
||||
// BFloat16 is always supported on RISC-V: natively when RISCV_BF16_SUPPORT
|
||||
// is defined, otherwise via the FP32-simulation fallback path.
|
||||
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
|
||||
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
|
||||
|
||||
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
|
||||
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
|
||||
@@ -486,9 +482,18 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
}
|
||||
|
||||
FP32Vec8 exp() const {
|
||||
// Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
|
||||
// Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
|
||||
// Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
|
||||
constexpr float exp_lo = -87.3365447505f; // ln(FLT_MIN)
|
||||
constexpr float exp_hi = 88.7228391117f; // ln(FLT_MAX)
|
||||
fixed_fp32x8_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_256)(
|
||||
RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
|
||||
VEC_ELEM_NUM);
|
||||
|
||||
const float inv_ln2 = 1.44269504088896341f;
|
||||
fixed_fp32x8_t x_scaled =
|
||||
RVVI(__riscv_vfmul_vf_f32, LMUL_256)(reg, inv_ln2, VEC_ELEM_NUM);
|
||||
RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x, inv_ln2, VEC_ELEM_NUM);
|
||||
fixed_i32x8_t n_int =
|
||||
RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_256)(x_scaled, VEC_ELEM_NUM);
|
||||
fixed_fp32x8_t n_float =
|
||||
@@ -706,9 +711,18 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
}
|
||||
|
||||
FP32Vec16 exp() const {
|
||||
// Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
|
||||
// Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
|
||||
// Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
|
||||
constexpr float exp_lo = -87.3365447505f; // ln(FLT_MIN)
|
||||
constexpr float exp_hi = 88.7228391117f; // ln(FLT_MAX)
|
||||
fixed_fp32x16_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_512)(
|
||||
RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
|
||||
VEC_ELEM_NUM);
|
||||
|
||||
const float inv_ln2 = 1.44269504088896341f;
|
||||
fixed_fp32x16_t x_scaled =
|
||||
RVVI(__riscv_vfmul_vf_f32, LMUL_512)(reg, inv_ln2, VEC_ELEM_NUM);
|
||||
RVVI(__riscv_vfmul_vf_f32, LMUL_512)(x, inv_ln2, VEC_ELEM_NUM);
|
||||
fixed_i32x16_t n_int =
|
||||
RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(x_scaled, VEC_ELEM_NUM);
|
||||
fixed_fp32x16_t n_float =
|
||||
|
||||
@@ -178,7 +178,12 @@ void rotary_embedding_gptj_impl(
|
||||
|
||||
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
|
||||
std::optional<torch::Tensor> key, int64_t head_size,
|
||||
torch::Tensor& cos_sin_cache, bool is_neox) {
|
||||
torch::Tensor& cos_sin_cache, bool is_neox,
|
||||
int64_t rope_dim_offset, bool inverse) {
|
||||
TORCH_CHECK(rope_dim_offset == 0,
|
||||
"rope_dim_offset != 0 is not supported on CPU");
|
||||
TORCH_CHECK(!inverse, "inverse rotary embedding is not supported on CPU");
|
||||
|
||||
int num_tokens = positions.numel();
|
||||
int rot_dim = cos_sin_cache.size(1);
|
||||
int num_heads = query.size(-1) / head_size;
|
||||
|
||||
@@ -263,7 +263,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def(
|
||||
"rotary_embedding(Tensor positions, Tensor! query,"
|
||||
" Tensor!? key, int head_size,"
|
||||
" Tensor cos_sin_cache, bool is_neox) -> ()");
|
||||
" Tensor cos_sin_cache, bool is_neox, int "
|
||||
"rope_dim_offset=0, bool inverse=False) -> ()");
|
||||
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
|
||||
|
||||
// Quantization
|
||||
|
||||
@@ -0,0 +1,477 @@
|
||||
/*
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
* SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
*
|
||||
* Horizontally-fused DeepseekV4-MLA kernel:
|
||||
* - Q side: per-head RMSNorm (no weight) + GPT-J RoPE on last ROPE_DIM
|
||||
* - KV side: GPT-J RoPE on last ROPE_DIM + UE8M0 FP8 quant on NoPE + paged
|
||||
* cache insert
|
||||
*
|
||||
* Structured after `applyMLARopeAndAssignQKVKernelGeneration` in
|
||||
* TensorRT-LLM's mlaKernels.cu: one kernel, one grid, with head-slot
|
||||
* dispatch choosing Q vs KV work per warp. The per-warp RMSNorm/RoPE
|
||||
* skeleton is adapted from vllm-deepseek_v4's existing
|
||||
* `fusedQKNormRopeKernel` (csrc/fused_qknorm_rope_kernel.cu).
|
||||
*
|
||||
* Assumptions (hard-coded for DeepseekV4 attention):
|
||||
* HEAD_DIM = 512
|
||||
* ROPE_DIM = 64 (RoPE applied to dims [NOPE_DIM, HEAD_DIM))
|
||||
* NOPE_DIM = 448
|
||||
* QUANT_BLOCK = 64 (UE8M0 FP8 quant block)
|
||||
* FP8_MAX = 448.0f
|
||||
* is_neox=false (GPT-J interleaved pairs)
|
||||
* cos_sin_cache layout [max_pos, rope_dim] = cos || sin (cos first, sin
|
||||
* second along last dim; each half is rope_dim/2 = 32 values)
|
||||
*
|
||||
* Cache layout per paged-cache block (block_size tokens):
|
||||
* [0, bs*576): token data, 448 fp8 + 128 bf16 each
|
||||
* [bs*576, bs*576 + bs*8): UE8M0 scales, 7 real + 1 pad per token
|
||||
*/
|
||||
|
||||
#include <cmath>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <type_traits>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/cuda.h>
|
||||
|
||||
#include "cuda_compat.h"
|
||||
#include "dispatch_utils.h"
|
||||
#include "type_convert.cuh"
|
||||
|
||||
#ifndef FINAL_MASK
|
||||
#define FINAL_MASK 0xffffffffu
|
||||
#endif
|
||||
|
||||
namespace vllm {
|
||||
namespace deepseek_v4_fused_ops {
|
||||
|
||||
namespace {
|
||||
inline int getSMVersion() {
|
||||
auto* props = at::cuda::getCurrentDeviceProperties();
|
||||
return props->major * 10 + props->minor;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Constants
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
constexpr int kHeadDim = 512;
|
||||
constexpr int kRopeDim = 64;
|
||||
constexpr int kNopeDim = kHeadDim - kRopeDim; // 448
|
||||
constexpr int kQuantBlock = 64;
|
||||
constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7
|
||||
constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad)
|
||||
constexpr int kTokenDataBytes = kNopeDim + kRopeDim * 2; // 448 + 128 = 576
|
||||
constexpr float kFp8Max = 448.0f;
|
||||
|
||||
// Per-warp layout: 32 lanes × 16 elems/lane = 512 elems = HEAD_DIM.
|
||||
constexpr int kNumLanes = 32;
|
||||
constexpr int kElemsPerLane = kHeadDim / kNumLanes; // 16
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Small inline helpers
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
__device__ __forceinline__ float warp4MaxAbs(float val) {
|
||||
// Reduce absolute max across 4 consecutive lanes (lane id & 3 group).
|
||||
float peer = __shfl_xor_sync(FINAL_MASK, val, 1);
|
||||
val = fmaxf(val, peer);
|
||||
peer = __shfl_xor_sync(FINAL_MASK, val, 2);
|
||||
val = fmaxf(val, peer);
|
||||
return val;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ float warpSum(float val) {
|
||||
#pragma unroll
|
||||
for (int mask = 16; mask > 0; mask >>= 1) {
|
||||
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
|
||||
}
|
||||
return val;
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Kernel
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
//
|
||||
// Grid: 1D, gridDim.x = ceil(num_tokens_full * (num_heads_q + 1) /
|
||||
// warps_per_block) Block: blockDim.x = 256 threads (8 warps per block) Each
|
||||
// warp handles one (token, head_slot) pair. head_slot < num_heads_q →
|
||||
// Q branch (RMSNorm + RoPE, in place) head_slot == num_heads_q → KV
|
||||
// branch (RoPE + UE8M0 quant + insert)
|
||||
//
|
||||
// With DP padding, q/kv/position_ids can have more rows than slot_mapping.
|
||||
// The Q branch covers all `num_tokens_full` rows (downstream attention uses
|
||||
// them). The KV branch only inserts the first `num_tokens_insert` tokens
|
||||
// (= slot_mapping length) into the paged cache.
|
||||
//
|
||||
template <typename scalar_t_in>
|
||||
__global__ void fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel(
|
||||
scalar_t_in* __restrict__ q_inout, // [N, H, 512] bf16, in place
|
||||
scalar_t_in const* __restrict__ kv_in, // [N, 512] bf16
|
||||
uint8_t* __restrict__ k_cache, // [num_blocks, block_stride]
|
||||
int64_t const* __restrict__ slot_mapping, // [num_tokens_insert] i64
|
||||
int64_t const* __restrict__ position_ids, // [N] i64
|
||||
float const* __restrict__ cos_sin_cache, // [max_pos, 64] fp32
|
||||
float const eps,
|
||||
int const num_tokens_full, // = q.size(0) = kv.size(0)
|
||||
int const num_tokens_insert, // = slot_mapping.size(0), ≤ num_tokens_full
|
||||
int const num_heads_q, // H
|
||||
int const cache_block_size, // tokens per paged-cache block
|
||||
int const kv_block_stride) { // bytes per paged-cache block
|
||||
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
|
||||
// BF16 _typeConvert specialization is unavailable on pre-Ampere. The
|
||||
// DeepseekV4 kernel only runs with bf16 inputs in practice, so compile a
|
||||
// no-op stub for sm_70/sm_75 to keep multi-arch builds happy.
|
||||
if constexpr (std::is_same_v<scalar_t_in, c10::BFloat16>) {
|
||||
return;
|
||||
} else {
|
||||
#endif
|
||||
using Converter = vllm::_typeConvert<scalar_t_in>;
|
||||
|
||||
int const warpsPerBlock = blockDim.x / 32;
|
||||
int const warpId = threadIdx.x / 32;
|
||||
int const laneId = threadIdx.x % 32;
|
||||
int const globalWarpIdx = blockIdx.x * warpsPerBlock + warpId;
|
||||
|
||||
int const total_slots_per_token = num_heads_q + 1;
|
||||
int const tokenIdx = globalWarpIdx / total_slots_per_token;
|
||||
int const slotIdx = globalWarpIdx % total_slots_per_token;
|
||||
if (tokenIdx >= num_tokens_full) return;
|
||||
|
||||
bool const isKV = (slotIdx == num_heads_q);
|
||||
// KV branch: skip DP-padded tokens (no slot reserved for them).
|
||||
if (isKV && tokenIdx >= num_tokens_insert) return;
|
||||
|
||||
// PDL: wait for predecessor kernel (upstream q/kv producer) to signal
|
||||
// before touching any global memory. No-op when PDL is not enabled on
|
||||
// the launch. The CUDA runtime wrapper emits the griddepcontrol.wait
|
||||
// PTX with the required memory clobber internally.
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
|
||||
// Dim range this lane owns within the 512-wide head.
|
||||
int const dim_base = laneId * kElemsPerLane; // in [0, 512) step 16
|
||||
|
||||
// ── Load 16 bf16 → 16 fp32 registers (one 16-byte + one 16-byte LDG) ────
|
||||
float elements[kElemsPerLane];
|
||||
float sumOfSquares = 0.0f;
|
||||
|
||||
scalar_t_in const* src_ptr;
|
||||
if (isKV) {
|
||||
src_ptr = kv_in + static_cast<int64_t>(tokenIdx) * kHeadDim + dim_base;
|
||||
} else {
|
||||
int64_t const q_row_offset =
|
||||
(static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
|
||||
dim_base;
|
||||
src_ptr = q_inout + q_row_offset;
|
||||
}
|
||||
|
||||
// Two 16-byte loads per thread (8 bf16 each). Use uint4 as the vector
|
||||
// type and bitcast to scalar_t_in packed pairs for conversion.
|
||||
uint4 v0 = *reinterpret_cast<uint4 const*>(src_ptr);
|
||||
uint4 v1 = *reinterpret_cast<uint4 const*>(src_ptr + 8);
|
||||
|
||||
{
|
||||
typename Converter::packed_hip_type const* p0 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type const*>(&v0);
|
||||
typename Converter::packed_hip_type const* p1 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type const*>(&v1);
|
||||
// Each packed_hip_type holds 2 bf16 → 4 packed = 8 elems per uint4.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float2 f2 = Converter::convert(p0[i]);
|
||||
elements[2 * i] = f2.x;
|
||||
elements[2 * i + 1] = f2.y;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float2 f2 = Converter::convert(p1[i]);
|
||||
elements[8 + 2 * i] = f2.x;
|
||||
elements[8 + 2 * i + 1] = f2.y;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Q branch: RMSNorm with no weight (has_weight=False) ─────────────────
|
||||
// Variance + rsqrt + multiply all in fp32, no intermediate bf16 round.
|
||||
// The downstream bf16 round only happens at the final store.
|
||||
if (!isKV) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
sumOfSquares += elements[i] * elements[i];
|
||||
}
|
||||
sumOfSquares = warpSum<float>(sumOfSquares);
|
||||
float const rms_rcp =
|
||||
rsqrtf(sumOfSquares / static_cast<float>(kHeadDim) + eps);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
elements[i] = elements[i] * rms_rcp;
|
||||
}
|
||||
}
|
||||
|
||||
// ── GPT-J RoPE on dims [NOPE_DIM, HEAD_DIM) ─────────────────────────────
|
||||
// All math in fp32. cos_sin_cache is loaded as fp32 (its native storage).
|
||||
bool const is_rope_lane = dim_base >= kNopeDim;
|
||||
if (is_rope_lane) {
|
||||
int64_t const pos = position_ids[tokenIdx];
|
||||
constexpr int kHalfRope = kRopeDim / 2; // 32
|
||||
float const* cos_ptr = cos_sin_cache + pos * kRopeDim;
|
||||
float const* sin_ptr = cos_ptr + kHalfRope;
|
||||
|
||||
int const rope_local_base = dim_base - kNopeDim; // in [0, 64) step 16
|
||||
#pragma unroll
|
||||
for (int p = 0; p < kElemsPerLane / 2; p++) {
|
||||
int const pair_dim = rope_local_base + 2 * p;
|
||||
int const half_idx = pair_dim / 2;
|
||||
float const cos_v = VLLM_LDG(cos_ptr + half_idx);
|
||||
float const sin_v = VLLM_LDG(sin_ptr + half_idx);
|
||||
float const x_even = elements[2 * p];
|
||||
float const x_odd = elements[2 * p + 1];
|
||||
elements[2 * p] = x_even * cos_v - x_odd * sin_v;
|
||||
elements[2 * p + 1] = x_even * sin_v + x_odd * cos_v;
|
||||
}
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// Q branch: cast to bf16 and store back in place.
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
if (!isKV) {
|
||||
uint4 out0, out1;
|
||||
typename Converter::packed_hip_type* po0 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
|
||||
typename Converter::packed_hip_type* po1 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po0[i] = Converter::convert(
|
||||
make_float2(elements[2 * i], elements[2 * i + 1]));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po1[i] = Converter::convert(
|
||||
make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
|
||||
}
|
||||
scalar_t_in* dst =
|
||||
q_inout +
|
||||
(static_cast<int64_t>(tokenIdx) * num_heads_q + slotIdx) * kHeadDim +
|
||||
dim_base;
|
||||
*reinterpret_cast<uint4*>(dst) = out0;
|
||||
*reinterpret_cast<uint4*>(dst + 8) = out1;
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
return;
|
||||
}
|
||||
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
// KV branch.
|
||||
// ═══════════════════════════════════════════════════════════════════════
|
||||
int64_t const slot_id = slot_mapping[tokenIdx];
|
||||
if (slot_id < 0) {
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
return;
|
||||
}
|
||||
|
||||
int64_t const block_idx = slot_id / cache_block_size;
|
||||
int64_t const pos_in_block = slot_id % cache_block_size;
|
||||
uint8_t* block_base =
|
||||
k_cache + block_idx * static_cast<int64_t>(kv_block_stride);
|
||||
uint8_t* token_fp8_ptr = block_base + pos_in_block * kTokenDataBytes;
|
||||
uint8_t* token_bf16_ptr = token_fp8_ptr + kNopeDim;
|
||||
uint8_t* token_scale_ptr =
|
||||
block_base + static_cast<int64_t>(cache_block_size) * kTokenDataBytes +
|
||||
pos_in_block * kScaleBytesPerToken;
|
||||
|
||||
// Round K to bf16 first, matching the unfused reference path where K is
|
||||
// materialized as bf16 before K quantization. absmax, clamp, and FP8
|
||||
// quant below all run on these bf16-rounded values.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
elements[i] = Converter::convert(Converter::convert(elements[i]));
|
||||
}
|
||||
|
||||
// Per-quant-block absmax must be computed by ALL 32 lanes (warp-collective
|
||||
// shuffle requires full participation). RoPE lanes contribute garbage,
|
||||
// but their values are gated out below via `!is_rope_lane`.
|
||||
float local_absmax = 0.0f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
local_absmax = fmaxf(local_absmax, fabsf(elements[i]));
|
||||
}
|
||||
float const absmax = fmaxf(warp4MaxAbs(local_absmax), 1e-4f);
|
||||
float const exponent = ceilf(log2f(absmax / kFp8Max));
|
||||
float const inv_scale = exp2f(-exponent);
|
||||
|
||||
if (!is_rope_lane) {
|
||||
// ── NoPE lane: UE8M0 FP8 quant ───────────────────────────────────────
|
||||
uint8_t out_bytes[kElemsPerLane];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < kElemsPerLane; i++) {
|
||||
float scaled = elements[i] * inv_scale;
|
||||
scaled = fminf(fmaxf(scaled, -kFp8Max), kFp8Max);
|
||||
__nv_fp8_storage_t s =
|
||||
__nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
|
||||
out_bytes[i] = static_cast<uint8_t>(s);
|
||||
}
|
||||
// One 16-byte STG per lane.
|
||||
*reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
|
||||
*reinterpret_cast<uint4 const*>(out_bytes);
|
||||
|
||||
// Lane (4k) of each 4-lane group writes the scale byte for block k<7.
|
||||
if ((laneId & 3) == 0) {
|
||||
int const q_block_idx = laneId >> 2; // 0..6 for NoPE lanes
|
||||
float encoded = fmaxf(fminf(exponent + 127.0f, 255.0f), 0.0f);
|
||||
token_scale_ptr[q_block_idx] = static_cast<uint8_t>(encoded);
|
||||
}
|
||||
// Lane 0 also writes the padding byte at index 7.
|
||||
if (laneId == 0) {
|
||||
token_scale_ptr[kNumQuantBlocks] = 0; // pad
|
||||
}
|
||||
} else {
|
||||
// ── RoPE lane: cast back to bf16 and store to cache bf16 tail ────────
|
||||
uint4 out0, out1;
|
||||
typename Converter::packed_hip_type* po0 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out0);
|
||||
typename Converter::packed_hip_type* po1 =
|
||||
reinterpret_cast<typename Converter::packed_hip_type*>(&out1);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po0[i] = Converter::convert(
|
||||
make_float2(elements[2 * i], elements[2 * i + 1]));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; i++) {
|
||||
po1[i] = Converter::convert(
|
||||
make_float2(elements[8 + 2 * i], elements[8 + 2 * i + 1]));
|
||||
}
|
||||
int const rope_local_base = dim_base - kNopeDim; // in [0, 64)
|
||||
scalar_t_in* bf16_dst =
|
||||
reinterpret_cast<scalar_t_in*>(token_bf16_ptr) + rope_local_base;
|
||||
*reinterpret_cast<uint4*>(bf16_dst) = out0;
|
||||
*reinterpret_cast<uint4*>(bf16_dst + 8) = out1;
|
||||
}
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) && !defined(USE_ROCM)
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Launch wrapper
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
template <typename scalar_t_in>
|
||||
void launchFusedDeepseekV4QNormRopeKVRopeQuantInsert(
|
||||
scalar_t_in* q_inout, scalar_t_in const* kv_in, uint8_t* k_cache,
|
||||
int64_t const* slot_mapping, int64_t const* position_ids,
|
||||
float const* cos_sin_cache, float const eps, int const num_tokens_full,
|
||||
int const num_tokens_insert, int const num_heads_q,
|
||||
int const cache_block_size, int const kv_block_stride,
|
||||
cudaStream_t stream) {
|
||||
constexpr int kBlockSize = 256;
|
||||
constexpr int kWarpsPerBlock = kBlockSize / 32;
|
||||
int64_t const total_warps =
|
||||
static_cast<int64_t>(num_tokens_full) * (num_heads_q + 1);
|
||||
int const grid =
|
||||
static_cast<int>((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock);
|
||||
|
||||
// PDL: enable programmatic stream serialization whenever the hardware
|
||||
// supports it (SM90+). On pre-Hopper GPUs the attribute is unavailable,
|
||||
// so leave numAttrs = 0 and launch as a regular kernel.
|
||||
static int const sm_version = getSMVersion();
|
||||
// Host-side guard: the device kernel body is compiled as a no-op for
|
||||
// bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
|
||||
// unavailable there. Refuse the launch loudly instead of silently
|
||||
// skipping the work.
|
||||
TORCH_CHECK(
|
||||
sm_version >= 80,
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert requires sm_80+ "
|
||||
"(Ampere or newer); got sm_",
|
||||
sm_version);
|
||||
cudaLaunchConfig_t config;
|
||||
config.gridDim = dim3(grid);
|
||||
config.blockDim = dim3(kBlockSize);
|
||||
config.dynamicSmemBytes = 0;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attrs[1];
|
||||
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attrs[0].val.programmaticStreamSerializationAllowed = 1;
|
||||
config.attrs = attrs;
|
||||
config.numAttrs = (sm_version >= 90) ? 1 : 0;
|
||||
|
||||
cudaLaunchKernelEx(
|
||||
&config, fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>,
|
||||
q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache, eps,
|
||||
num_tokens_full, num_tokens_insert, num_heads_q, cache_block_size,
|
||||
kv_block_stride);
|
||||
}
|
||||
|
||||
} // namespace deepseek_v4_fused_ops
|
||||
} // namespace vllm
|
||||
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
// Torch op wrapper
|
||||
// ────────────────────────────────────────────────────────────────────────────
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
|
||||
torch::Tensor& q, // [N, H, 512] bf16, in place
|
||||
torch::Tensor const& kv, // [N, 512] bf16 (read-only)
|
||||
torch::Tensor& k_cache, // [num_blocks, block_bytes] uint8
|
||||
torch::Tensor const& slot_mapping, // [N] int64
|
||||
torch::Tensor const& position_ids, // [N] int64
|
||||
torch::Tensor const& cos_sin_cache, // [max_pos, rope_dim] bf16
|
||||
double eps, int64_t cache_block_size) {
|
||||
TORCH_CHECK(q.is_cuda() && q.is_contiguous(), "q must be contiguous CUDA");
|
||||
TORCH_CHECK(kv.is_cuda() && kv.is_contiguous(), "kv must be contiguous CUDA");
|
||||
TORCH_CHECK(k_cache.is_cuda(), "k_cache must be CUDA");
|
||||
TORCH_CHECK(slot_mapping.is_cuda() && slot_mapping.dtype() == torch::kInt64,
|
||||
"slot_mapping must be int64 CUDA");
|
||||
TORCH_CHECK(position_ids.is_cuda() && position_ids.dtype() == torch::kInt64,
|
||||
"position_ids must be int64 CUDA");
|
||||
TORCH_CHECK(cos_sin_cache.is_cuda(), "cos_sin_cache must be CUDA");
|
||||
TORCH_CHECK(q.dim() == 3 && q.size(2) == 512, "q shape [N, H, 512]");
|
||||
TORCH_CHECK(kv.dim() == 2 && kv.size(1) == 512, "kv shape [N, 512]");
|
||||
TORCH_CHECK(q.dtype() == kv.dtype(), "q and kv dtype must match");
|
||||
TORCH_CHECK(k_cache.dtype() == torch::kUInt8, "k_cache must be uint8");
|
||||
TORCH_CHECK(cos_sin_cache.dim() == 2 && cos_sin_cache.size(1) == 64,
|
||||
"cos_sin_cache shape [max_pos, 64]");
|
||||
TORCH_CHECK(cos_sin_cache.dtype() == torch::kFloat32,
|
||||
"cos_sin_cache must be float32");
|
||||
|
||||
// With DP padding, slot_mapping can be shorter than q/kv/positions.
|
||||
// Q-norm+RoPE runs on all q.size(0) rows (downstream attention uses them);
|
||||
// KV quant+insert runs only on the first slot_mapping.size(0) rows.
|
||||
int const num_tokens_full = static_cast<int>(q.size(0));
|
||||
int const num_tokens_insert = static_cast<int>(slot_mapping.size(0));
|
||||
TORCH_CHECK(static_cast<int>(kv.size(0)) == num_tokens_full &&
|
||||
static_cast<int>(position_ids.size(0)) == num_tokens_full,
|
||||
"q/kv/position_ids row counts must match");
|
||||
TORCH_CHECK(num_tokens_insert <= num_tokens_full,
|
||||
"slot_mapping must not exceed q row count");
|
||||
int const num_heads_q = static_cast<int>(q.size(1));
|
||||
int const cache_block_size_i = static_cast<int>(cache_block_size);
|
||||
int const kv_block_stride = static_cast<int>(k_cache.stride(0));
|
||||
|
||||
at::cuda::OptionalCUDAGuard device_guard(device_of(q));
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
VLLM_DISPATCH_HALF_TYPES(
|
||||
q.scalar_type(), "fused_deepseek_v4_qnorm_rope_kv_insert", [&] {
|
||||
using qkv_scalar_t = scalar_t;
|
||||
vllm::deepseek_v4_fused_ops::
|
||||
launchFusedDeepseekV4QNormRopeKVRopeQuantInsert<qkv_scalar_t>(
|
||||
reinterpret_cast<qkv_scalar_t*>(q.data_ptr()),
|
||||
reinterpret_cast<qkv_scalar_t const*>(kv.data_ptr()),
|
||||
reinterpret_cast<uint8_t*>(k_cache.data_ptr()),
|
||||
reinterpret_cast<int64_t const*>(slot_mapping.data_ptr()),
|
||||
reinterpret_cast<int64_t const*>(position_ids.data_ptr()),
|
||||
cos_sin_cache.data_ptr<float>(), static_cast<float>(eps),
|
||||
num_tokens_full, num_tokens_insert, num_heads_q,
|
||||
cache_block_size_i, kv_block_stride, stream);
|
||||
});
|
||||
}
|
||||
@@ -77,7 +77,8 @@ __global__ void rms_norm_kernel(
|
||||
#pragma unroll
|
||||
for (int j = 0; j < VEC_SIZE; j++) {
|
||||
float x = static_cast<float>(src1.val[j]);
|
||||
dst.val[j] = ((scalar_t)(x * s_variance)) * src2.val[j];
|
||||
float w = static_cast<float>(src2.val[j]);
|
||||
dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
|
||||
}
|
||||
v_out[i] = dst;
|
||||
}
|
||||
@@ -134,10 +135,17 @@ fused_add_rms_norm_kernel(
|
||||
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
|
||||
int id = blockIdx.x * vec_hidden_size + idx;
|
||||
int64_t strided_id = blockIdx.x * vec_input_stride + idx;
|
||||
_f16Vec<scalar_t, width> temp = residual_v[id];
|
||||
temp *= s_variance;
|
||||
temp *= weight_v[idx];
|
||||
input_v[strided_id] = temp;
|
||||
_f16Vec<scalar_t, width> res = residual_v[id];
|
||||
_f16Vec<scalar_t, width> w = weight_v[idx];
|
||||
_f16Vec<scalar_t, width> out;
|
||||
using Converter = _typeConvert<scalar_t>;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < width; ++j) {
|
||||
float x = Converter::convert(res.data[j]);
|
||||
float wf = Converter::convert(w.data[j]);
|
||||
out.data[j] = Converter::convert(x * s_variance * wf);
|
||||
}
|
||||
input_v[strided_id] = out;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -174,8 +182,8 @@ fused_add_rms_norm_kernel(
|
||||
|
||||
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
|
||||
float x = (float)residual[blockIdx.x * hidden_size + idx];
|
||||
input[blockIdx.x * input_stride + idx] =
|
||||
((scalar_t)(x * s_variance)) * weight[idx];
|
||||
float w = (float)weight[idx];
|
||||
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -65,9 +65,16 @@ __global__ void rms_norm_static_fp8_quant_kernel(
|
||||
#pragma unroll
|
||||
for (int j = 0; j < VEC_SIZE; j++) {
|
||||
float x = static_cast<float>(src1.val[j]);
|
||||
float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j];
|
||||
float w = static_cast<float>(src2.val[j]);
|
||||
// Round normalized result through scalar_t to match the precision of the
|
||||
// unfused composite (rms_norm writes scalar_t, then
|
||||
// static_scaled_fp8_quant re-loads it as float before FP8 conversion).
|
||||
// Without this round, the fused path is strictly more accurate and
|
||||
// disagrees with the composite at exact E4M3 quantization tie boundaries.
|
||||
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
|
||||
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
|
||||
scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
|
||||
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
|
||||
scale_inv);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -127,13 +134,21 @@ fused_add_rms_norm_static_fp8_quant_kernel(
|
||||
|
||||
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
|
||||
int id = blockIdx.x * vec_hidden_size + idx;
|
||||
_f16Vec<scalar_t, width> temp = residual_v[id];
|
||||
temp *= s_variance;
|
||||
temp *= weight_v[idx];
|
||||
_f16Vec<scalar_t, width> res = residual_v[id];
|
||||
_f16Vec<scalar_t, width> w = weight_v[idx];
|
||||
using Converter = _typeConvert<scalar_t>;
|
||||
using HipT = typename Converter::hip_type;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < width; ++i) {
|
||||
out[id * width + i] =
|
||||
scaled_fp8_conversion<true, fp8_type>(float(temp.data[i]), scale_inv);
|
||||
float x = Converter::convert(res.data[i]);
|
||||
float wf = Converter::convert(w.data[i]);
|
||||
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
|
||||
// to match the unfused composite path at FP8 boundaries. We use the
|
||||
// backend's hip_type for the intermediate since c10::Half/BFloat16 has
|
||||
// ambiguous conversions on CUDA and no implicit conversion on ROCm.
|
||||
HipT out_norm_h = Converter::convert(x * s_variance * wf);
|
||||
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
|
||||
Converter::convert(out_norm_h), scale_inv);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -176,9 +191,12 @@ fused_add_rms_norm_static_fp8_quant_kernel(
|
||||
|
||||
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
|
||||
float x = (float)residual[blockIdx.x * hidden_size + idx];
|
||||
float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
|
||||
out[blockIdx.x * hidden_size + idx] =
|
||||
scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
|
||||
float w = (float)weight[idx];
|
||||
// See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
|
||||
// to match the unfused composite path at FP8 boundaries.
|
||||
scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
|
||||
out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
|
||||
static_cast<float>(out_norm), scale_inv);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -277,7 +277,9 @@ void quant_impl(void* output, void* output_scale, void* input,
|
||||
(totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
|
||||
if (blockRepeat > 1) {
|
||||
size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
|
||||
if (n_experts >= 4) {
|
||||
// The shared-memory vectorized offset load only handles full 4-expert
|
||||
// chunks. Use the scalar specialization for the remainder cases.
|
||||
if (n_experts >= 4 && n_experts % 4 == 0) {
|
||||
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
|
||||
<<<grid, block, shared_mem_size, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
@@ -299,7 +301,9 @@ void quant_impl(void* output, void* output_scale, void* input,
|
||||
n_experts);
|
||||
}
|
||||
} else {
|
||||
if (n_experts >= 16) {
|
||||
// The low-latency vectorized expert lookup only handles full 16-expert
|
||||
// chunks. Fall back to the scalar lookup path for the remainder cases.
|
||||
if (n_experts >= 16 && n_experts % 16 == 0) {
|
||||
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
|
||||
<<<grid, block, 0, stream>>>(
|
||||
m_topk, k, reinterpret_cast<T*>(input),
|
||||
|
||||
@@ -12,6 +12,15 @@ void topk_sigmoid(torch::Tensor& topk_weights, torch::Tensor& topk_indices,
|
||||
torch::Tensor& gating_output, bool renormalize,
|
||||
std::optional<torch::Tensor> bias);
|
||||
|
||||
void topk_softplus_sqrt(torch::Tensor& topk_weights,
|
||||
torch::Tensor& topk_indices,
|
||||
torch::Tensor& token_expert_indices,
|
||||
torch::Tensor& gating_output, bool renormalize,
|
||||
double routed_scaling_factor,
|
||||
const c10::optional<torch::Tensor>& correction_bias,
|
||||
const c10::optional<torch::Tensor>& input_ids,
|
||||
const c10::optional<torch::Tensor>& tid2eid);
|
||||
|
||||
void moe_sum(torch::Tensor& input, torch::Tensor& output);
|
||||
|
||||
void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
|
||||
|
||||
@@ -0,0 +1,715 @@
|
||||
/*
|
||||
* Adapted from
|
||||
* https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
|
||||
* Copyright (c) 2024, The vLLM team.
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
|
||||
* AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#include <type_traits>
|
||||
#include <torch/all.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include "../cuda_compat.h"
|
||||
#include "../cub_helpers.h"
|
||||
#ifndef USE_ROCM
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp16.h>
|
||||
#else
|
||||
#include <hip/hip_bf16.h>
|
||||
#include <hip/hip_fp16.h>
|
||||
typedef __hip_bfloat16 __nv_bfloat16;
|
||||
typedef __hip_bfloat162 __nv_bfloat162;
|
||||
#endif
|
||||
|
||||
#define MAX(a, b) ((a) > (b) ? (a) : (b))
|
||||
#define MIN(a, b) ((a) < (b) ? (a) : (b))
|
||||
|
||||
namespace vllm {
|
||||
namespace moe {
|
||||
|
||||
/// Aligned array type
|
||||
template <typename T,
|
||||
/// Number of elements in the array
|
||||
int N,
|
||||
/// Alignment requirement in bytes
|
||||
int Alignment = sizeof(T) * N>
|
||||
struct alignas(Alignment) AlignedArray {
|
||||
T data[N];
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ float toFloat(T value) {
|
||||
if constexpr (std::is_same_v<T, float>) {
|
||||
return value;
|
||||
} else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
|
||||
return __bfloat162float(value);
|
||||
} else if constexpr (std::is_same_v<T, __half>) {
|
||||
return __half2float(value);
|
||||
}
|
||||
}
|
||||
|
||||
#define FINAL_MASK 0xffffffff
|
||||
template <typename T>
|
||||
__inline__ __device__ T warpReduceSum(T val) {
|
||||
#pragma unroll
|
||||
for (int mask = 16; mask > 0; mask >>= 1)
|
||||
val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
|
||||
return val;
|
||||
}
|
||||
|
||||
// ====================== TopK softplus_sqrt things
|
||||
// ===============================
|
||||
|
||||
/*
|
||||
A Top-K gating softplus_sqrt written to exploit when the number of experts in
|
||||
the MoE layers are a small power of 2. This allows us to cleanly share the
|
||||
rows among the threads in a single warp and eliminate communication between
|
||||
warps (so no need to use shared mem).
|
||||
|
||||
It fuses the sigmoid, max and argmax into a single kernel.
|
||||
|
||||
Limitations:
|
||||
1) This implementation is optimized for when the number of experts is a small
|
||||
power of 2. Additionally it also supports when number of experts is multiple
|
||||
of 64 which is still faster than the computing sigmoid and topK separately
|
||||
(only tested on CUDA yet). 2) This implementation assumes k is small, but will
|
||||
work for any k.
|
||||
*/
|
||||
|
||||
template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG,
|
||||
int WARP_SIZE_PARAM, bool USE_HASH, typename IndType,
|
||||
typename InputType = float>
|
||||
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
void topkGatingSoftplusSqrt(
|
||||
const InputType* input, const bool* finished, float* output,
|
||||
const int num_rows, IndType* indices, int* source_rows, const int k,
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const IndType* input_ids, const IndType* tid2eid) {
|
||||
static_assert(std::is_same_v<InputType, float> ||
|
||||
std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>,
|
||||
"InputType must be float, __nv_bfloat16, or __half");
|
||||
|
||||
// We begin by enforcing compile time assertions and setting up compile time
|
||||
// constants.
|
||||
static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
|
||||
"BYTES_PER_LDG must be power of 2");
|
||||
static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
|
||||
|
||||
// Number of bytes each thread pulls in per load
|
||||
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
|
||||
static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
|
||||
static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
|
||||
static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
|
||||
|
||||
if constexpr (std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>) {
|
||||
static_assert(ELTS_PER_LDG == 1 || ELTS_PER_LDG % 2 == 0,
|
||||
"ELTS_PER_LDG must be 1 or even for 16-bit conversion");
|
||||
}
|
||||
|
||||
// Restrictions based on previous section.
|
||||
static_assert(
|
||||
VPT % ELTS_PER_LDG == 0,
|
||||
"The elements per thread must be a multiple of the elements per ldg");
|
||||
static_assert(WARP_SIZE_PARAM % THREADS_PER_ROW == 0,
|
||||
"The threads per row must cleanly divide the threads per warp");
|
||||
static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
|
||||
"THREADS_PER_ROW must be power of 2");
|
||||
static_assert(THREADS_PER_ROW <= WARP_SIZE_PARAM,
|
||||
"THREADS_PER_ROW can be at most warp size");
|
||||
|
||||
// We have NUM_EXPERTS elements per row. We specialize for small #experts
|
||||
static constexpr int ELTS_PER_WARP = WARP_SIZE_PARAM * VPT;
|
||||
static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
|
||||
static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
|
||||
|
||||
// Restrictions for previous section.
|
||||
static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
|
||||
"The elts per row must cleanly divide the total elt per warp");
|
||||
|
||||
// ===================== From this point, we finally start computing run-time
|
||||
// variables. ========================
|
||||
|
||||
// Compute CTA and warp rows. We pack multiple rows into a single warp, and a
|
||||
// block contains WARPS_PER_CTA warps. This, each block processes a chunk of
|
||||
// rows. We start by computing the start row for each block.
|
||||
const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
|
||||
|
||||
// Now, using the base row per thread block, we compute the base row per warp.
|
||||
const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
|
||||
|
||||
// The threads in a warp are split into sub-groups that will work on a row.
|
||||
// We compute row offset for each thread sub-group
|
||||
const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
|
||||
const int thread_row = warp_base_row + thread_row_in_warp;
|
||||
|
||||
// Threads with indices out of bounds should early exit here.
|
||||
if (thread_row >= num_rows) {
|
||||
return;
|
||||
}
|
||||
const bool row_is_active = finished ? !finished[thread_row] : true;
|
||||
|
||||
// We finally start setting up the read pointers for each thread. First, each
|
||||
// thread jumps to the start of the row it will read.
|
||||
const InputType* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
|
||||
|
||||
// Now, we compute the group each thread belong to in order to determine the
|
||||
// first column to start loads.
|
||||
const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
|
||||
const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
|
||||
const InputType* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
|
||||
|
||||
// Finally, we pull in the data from global mem
|
||||
float row_chunk[VPT];
|
||||
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#endif
|
||||
|
||||
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
|
||||
// to float
|
||||
if constexpr (std::is_same_v<InputType, float>) {
|
||||
using VecType = AlignedArray<float, ELTS_PER_LDG>;
|
||||
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
|
||||
const VecType* vec_thread_read_ptr =
|
||||
reinterpret_cast<const VecType*>(thread_read_ptr);
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
|
||||
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
|
||||
}
|
||||
} else if constexpr (std::is_same_v<InputType, __nv_bfloat16>) {
|
||||
if constexpr (ELTS_PER_LDG >= 2) {
|
||||
using VecType = AlignedArray<__nv_bfloat16, ELTS_PER_LDG>;
|
||||
float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
|
||||
const VecType* vec_thread_read_ptr =
|
||||
reinterpret_cast<const VecType*>(thread_read_ptr);
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
|
||||
VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
|
||||
int base_idx_f2 = ii * ELTS_PER_LDG / 2;
|
||||
#pragma unroll
|
||||
for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
|
||||
row_chunk_f2[base_idx_f2 + jj] = __bfloat1622float2(
|
||||
*reinterpret_cast<const __nv_bfloat162*>(vec.data + jj * 2));
|
||||
}
|
||||
}
|
||||
} else { // ELTS_PER_LDG == 1
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
|
||||
const __nv_bfloat16* scalar_ptr =
|
||||
thread_read_ptr + ii * THREADS_PER_ROW;
|
||||
row_chunk[ii] = __bfloat162float(*scalar_ptr);
|
||||
}
|
||||
}
|
||||
} else if constexpr (std::is_same_v<InputType, __half>) {
|
||||
if constexpr (ELTS_PER_LDG >= 2) {
|
||||
using VecType = AlignedArray<__half, ELTS_PER_LDG>;
|
||||
float2* row_chunk_f2 = reinterpret_cast<float2*>(row_chunk);
|
||||
const VecType* vec_thread_read_ptr =
|
||||
reinterpret_cast<const VecType*>(thread_read_ptr);
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
|
||||
VecType vec = vec_thread_read_ptr[ii * THREADS_PER_ROW];
|
||||
int base_idx_f2 = ii * ELTS_PER_LDG / 2;
|
||||
#pragma unroll
|
||||
for (int jj = 0; jj < ELTS_PER_LDG / 2; ++jj) {
|
||||
row_chunk_f2[base_idx_f2 + jj] = __half22float2(
|
||||
*reinterpret_cast<const __half2*>(vec.data + jj * 2));
|
||||
}
|
||||
}
|
||||
} else { // ELTS_PER_LDG == 1
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
|
||||
const __half* scalar_ptr = thread_read_ptr + ii * THREADS_PER_ROW;
|
||||
row_chunk[ii] = __half2float(*scalar_ptr);
|
||||
}
|
||||
}
|
||||
}
|
||||
constexpr float threshold = 20.0f;
|
||||
constexpr float beta = 1.0f;
|
||||
|
||||
// Hash MoE path: indices are predetermined from lookup table
|
||||
if constexpr (USE_HASH) {
|
||||
const IndType token_id = input_ids[thread_row];
|
||||
const IndType* expert_indices_for_token = tid2eid + token_id * k;
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
row_chunk[ii] = sqrtf(val);
|
||||
}
|
||||
float selected_sum = 0.f;
|
||||
#pragma unroll
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int expert = expert_indices_for_token[k_idx];
|
||||
const int idx = k * thread_row + k_idx;
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
if (expert == expert_idx) {
|
||||
indices[idx] = expert;
|
||||
selected_sum += row_chunk[ii];
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
// Compute per-thread scale (using warp reduction when renormalizing).
|
||||
if (renormalize) {
|
||||
selected_sum = warpReduceSum(selected_sum);
|
||||
}
|
||||
float scale = static_cast<float>(routed_scaling_factor);
|
||||
if (renormalize) {
|
||||
const float denom = selected_sum > 0.f ? selected_sum : 1.f;
|
||||
scale /= denom;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int expert = expert_indices_for_token[k_idx];
|
||||
const int idx = k * thread_row + k_idx;
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
if (expert == expert_idx) {
|
||||
output[idx] = row_chunk[ii] * scale;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
return;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
// Compute softplus: log(1 + exp(val)) with numerical stability
|
||||
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
val = sqrtf(val);
|
||||
if (correction_bias) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
val = val + correction_bias[expert_idx];
|
||||
}
|
||||
row_chunk[ii] = val;
|
||||
}
|
||||
|
||||
// Original TopK path: find top-k experts by score
|
||||
// Now, sigmoid_res contains the sigmoid of the row chunk. Now, I want to find
|
||||
// the topk elements in each row, along with the max index.
|
||||
int start_col = first_elt_read_by_thread;
|
||||
static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
|
||||
|
||||
float selected_sum = 0.f;
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
// First, each thread does the local argmax
|
||||
float max_val = row_chunk[0];
|
||||
int expert = start_col;
|
||||
#pragma unroll
|
||||
for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
|
||||
++ldg, col += COLS_PER_GROUP_LDG) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
|
||||
float val = row_chunk[ldg * ELTS_PER_LDG + ii];
|
||||
|
||||
// No check on the experts here since columns with the smallest index
|
||||
// are processed first and only updated if > (not >=)
|
||||
if (val > max_val) {
|
||||
max_val = val;
|
||||
expert = col + ii;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Now, we perform the argmax reduce. We use the butterfly pattern so threads
|
||||
// reach consensus about the max. This will be useful for K > 1 so that the
|
||||
// threads can agree on "who" had the max value. That thread can then blank out
|
||||
// their max with -inf and the warp can run more iterations...
|
||||
#pragma unroll
|
||||
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
|
||||
float other_max =
|
||||
VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
|
||||
int other_expert =
|
||||
VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
|
||||
|
||||
// We want lower indices to "win" in every thread so we break ties this
|
||||
// way
|
||||
if (other_max > max_val ||
|
||||
(other_max == max_val && other_expert < expert)) {
|
||||
max_val = other_max;
|
||||
expert = other_expert;
|
||||
}
|
||||
}
|
||||
|
||||
// Write the max for this k iteration to global memory.
|
||||
if (thread_group_idx == 0) {
|
||||
// Add a guard to ignore experts not included by this node
|
||||
const bool node_uses_expert =
|
||||
expert >= start_expert && expert < end_expert;
|
||||
const bool should_process_row = row_is_active && node_uses_expert;
|
||||
|
||||
// The lead thread from each sub-group will write out the final results to
|
||||
// global memory. (This will be a single) thread per row of the
|
||||
// input/output matrices.
|
||||
const int idx = k * thread_row + k_idx;
|
||||
if (correction_bias != nullptr) {
|
||||
max_val -= correction_bias[expert];
|
||||
}
|
||||
output[idx] = max_val;
|
||||
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
|
||||
source_rows[idx] = k_idx * num_rows + thread_row;
|
||||
if (renormalize) {
|
||||
selected_sum += max_val;
|
||||
}
|
||||
}
|
||||
|
||||
// Finally, we clear the value in the thread with the current max if there
|
||||
// is another iteration to run.
|
||||
if (k_idx + 1 < k) {
|
||||
const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
|
||||
const int thread_to_clear_in_group =
|
||||
(expert / ELTS_PER_LDG) % THREADS_PER_ROW;
|
||||
|
||||
// Only the thread in the group which produced the max will reset the
|
||||
// "winning" value to -inf.
|
||||
if (thread_group_idx == thread_to_clear_in_group) {
|
||||
const int offset_for_expert = expert % ELTS_PER_LDG;
|
||||
// Safe to set to any negative value since row_chunk values must be
|
||||
// between 0 and 1.
|
||||
row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
|
||||
-10000.f;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Apply renormalization and routed scaling factor to final weights.
|
||||
if (thread_group_idx == 0) {
|
||||
float scale = static_cast<float>(routed_scaling_factor);
|
||||
if (renormalize) {
|
||||
const float denom = selected_sum > 0.f ? selected_sum : 1.f;
|
||||
scale /= denom;
|
||||
}
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int idx = k * thread_row + k_idx;
|
||||
output[idx] = output[idx] * scale;
|
||||
}
|
||||
}
|
||||
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
|
||||
namespace detail {
|
||||
// Constructs some constants needed to partition the work across threads at
|
||||
// compile time.
|
||||
template <int EXPERTS, int BYTES_PER_LDG, int WARP_SIZE_PARAM,
|
||||
typename InputType>
|
||||
struct TopkConstants {
|
||||
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(InputType);
|
||||
static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0 ||
|
||||
EXPERTS % (ELTS_PER_LDG * WARP_SIZE_PARAM) == 0,
|
||||
"");
|
||||
static constexpr int VECs_PER_THREAD =
|
||||
MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE_PARAM));
|
||||
static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
|
||||
static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
|
||||
static const int ROWS_PER_WARP = WARP_SIZE_PARAM / THREADS_PER_ROW;
|
||||
};
|
||||
} // namespace detail
|
||||
|
||||
#define DISPATCH_HASH(use_hash, USE_HASH, ...) \
|
||||
if (use_hash) { \
|
||||
const bool USE_HASH = true; \
|
||||
static_assert(USE_HASH == true, "USE_HASH must be compile-time constant"); \
|
||||
__VA_ARGS__ \
|
||||
} else { \
|
||||
const bool USE_HASH = false; \
|
||||
static_assert(USE_HASH == false, \
|
||||
"USE_HASH must be compile-time constant"); \
|
||||
__VA_ARGS__ \
|
||||
}
|
||||
|
||||
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM,
|
||||
int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
|
||||
void topkGatingSoftplusSqrtLauncherHelper(
|
||||
const InputType* input, const bool* finished, float* output,
|
||||
IndType* indices, int* source_row, const int num_rows, const int k,
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const bool use_hash, const IndType* input_ids, const IndType* tid2eid,
|
||||
cudaStream_t stream) {
|
||||
static constexpr int BYTES_PER_LDG =
|
||||
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
|
||||
using Constants =
|
||||
detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
|
||||
static constexpr int VPT = Constants::VPT;
|
||||
static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
|
||||
const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
|
||||
const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
|
||||
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
|
||||
DISPATCH_HASH(use_hash, USE_HASH, {
|
||||
auto* kernel =
|
||||
&topkGatingSoftplusSqrt<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG,
|
||||
WARP_SIZE_PARAM, USE_HASH, IndType, InputType>;
|
||||
#ifndef USE_ROCM
|
||||
cudaLaunchConfig_t config = {};
|
||||
config.gridDim = num_blocks;
|
||||
config.blockDim = block_dim;
|
||||
config.dynamicSmemBytes = 0;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attrs[1];
|
||||
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attrs[0].val.programmaticStreamSerializationAllowed = 1;
|
||||
config.numAttrs = 1;
|
||||
config.attrs = attrs;
|
||||
cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
|
||||
indices, source_row, k, start_expert, end_expert,
|
||||
renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid);
|
||||
#else
|
||||
kernel<<<num_blocks, block_dim, 0, stream>>>(
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert,
|
||||
end_expert, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid);
|
||||
#endif
|
||||
})
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
|
||||
static_assert(WARP_SIZE == 32, \
|
||||
"Unsupported warp size. Only 32 is supported for CUDA"); \
|
||||
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, WARP_SIZE, \
|
||||
MAX_BYTES>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
|
||||
stream);
|
||||
#else
|
||||
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
|
||||
if (WARP_SIZE == 64) { \
|
||||
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 64, \
|
||||
MAX_BYTES>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, \
|
||||
tid2eid, stream); \
|
||||
} else if (WARP_SIZE == 32) { \
|
||||
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
|
||||
MAX_BYTES>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, \
|
||||
tid2eid, stream); \
|
||||
} else { \
|
||||
assert(false && \
|
||||
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
|
||||
}
|
||||
#endif
|
||||
|
||||
template <typename IndType, typename InputType>
|
||||
void topkGatingSoftplusSqrtKernelLauncher(
|
||||
const InputType* gating_output, float* topk_weights, IndType* topk_indices,
|
||||
int* token_expert_indices, const int num_tokens, const int num_experts,
|
||||
const int topk, const bool renormalize, double routed_scaling_factor,
|
||||
const float* correction_bias, const bool use_hash, const IndType* input_ids,
|
||||
const IndType* tid2eid, cudaStream_t stream) {
|
||||
static constexpr int WARPS_PER_TB = 4;
|
||||
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
|
||||
#ifndef USE_ROCM
|
||||
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
|
||||
// elements can be loaded by a warp
|
||||
static constexpr int BYTES_PER_LDG_MULTIPLE_64 =
|
||||
(std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>)
|
||||
? 4
|
||||
: 8;
|
||||
#endif
|
||||
switch (num_experts) {
|
||||
case 1:
|
||||
LAUNCH_SOFTPLUS_SQRT(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 2:
|
||||
LAUNCH_SOFTPLUS_SQRT(2, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 4:
|
||||
LAUNCH_SOFTPLUS_SQRT(4, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 8:
|
||||
LAUNCH_SOFTPLUS_SQRT(8, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 16:
|
||||
LAUNCH_SOFTPLUS_SQRT(16, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 32:
|
||||
LAUNCH_SOFTPLUS_SQRT(32, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 64:
|
||||
LAUNCH_SOFTPLUS_SQRT(64, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 128:
|
||||
LAUNCH_SOFTPLUS_SQRT(128, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 256:
|
||||
LAUNCH_SOFTPLUS_SQRT(256, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
case 512:
|
||||
LAUNCH_SOFTPLUS_SQRT(512, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2);
|
||||
break;
|
||||
// (CUDA only) support multiples of 64 when num_experts is not power of 2.
|
||||
// ROCm uses WARP_SIZE 64 so 8 bytes loading won't fit for some of
|
||||
// num_experts, alternatively we can test 4 bytes loading and enable it in
|
||||
// future.
|
||||
#ifndef USE_ROCM
|
||||
case 192:
|
||||
LAUNCH_SOFTPLUS_SQRT(192, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
break;
|
||||
case 320:
|
||||
LAUNCH_SOFTPLUS_SQRT(320, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
break;
|
||||
case 384:
|
||||
LAUNCH_SOFTPLUS_SQRT(384, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
break;
|
||||
case 448:
|
||||
LAUNCH_SOFTPLUS_SQRT(448, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
break;
|
||||
case 576:
|
||||
LAUNCH_SOFTPLUS_SQRT(576, WARPS_PER_TB, BYTES_PER_LDG_MULTIPLE_64);
|
||||
break;
|
||||
#endif
|
||||
default: {
|
||||
TORCH_CHECK(false, "Unsupported expert number: ", num_experts);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace moe
|
||||
} // namespace vllm
|
||||
|
||||
template <typename ComputeType>
|
||||
void dispatch_topk_softplus_sqrt_launch(
|
||||
const ComputeType* gating_output, torch::Tensor& topk_weights,
|
||||
torch::Tensor& topk_indices, torch::Tensor& token_expert_indices,
|
||||
int num_tokens, int num_experts, int topk, bool renormalize,
|
||||
double routed_scaling_factor,
|
||||
const c10::optional<torch::Tensor>& correction_bias,
|
||||
const c10::optional<torch::Tensor>& input_ids,
|
||||
const c10::optional<torch::Tensor>& tid2eid, cudaStream_t stream) {
|
||||
const float* bias_ptr = nullptr;
|
||||
if (correction_bias.has_value()) {
|
||||
bias_ptr = correction_bias.value().data_ptr<float>();
|
||||
}
|
||||
bool use_hash = false;
|
||||
if (tid2eid.has_value()) {
|
||||
TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE");
|
||||
use_hash = true;
|
||||
}
|
||||
if (topk_indices.scalar_type() == at::ScalarType::Int) {
|
||||
const int* input_ids_ptr = nullptr;
|
||||
const int* tid2eid_ptr = nullptr;
|
||||
if (tid2eid.has_value()) {
|
||||
input_ids_ptr = input_ids.value().data_ptr<int>();
|
||||
tid2eid_ptr = tid2eid.value().data_ptr<int>();
|
||||
}
|
||||
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
|
||||
gating_output, topk_weights.data_ptr<float>(),
|
||||
topk_indices.data_ptr<int>(), token_expert_indices.data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
|
||||
} else if (topk_indices.scalar_type() == at::ScalarType::UInt32) {
|
||||
const uint32_t* input_ids_ptr = nullptr;
|
||||
const uint32_t* tid2eid_ptr = nullptr;
|
||||
if (tid2eid.has_value()) {
|
||||
input_ids_ptr = input_ids.value().data_ptr<uint32_t>();
|
||||
tid2eid_ptr = tid2eid.value().data_ptr<uint32_t>();
|
||||
}
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
|
||||
gating_output, topk_weights.data_ptr<float>(),
|
||||
topk_indices.data_ptr<uint32_t>(), token_expert_indices.data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
|
||||
} else {
|
||||
TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
|
||||
|
||||
const int64_t* input_ids_ptr = nullptr;
|
||||
const int64_t* tid2eid_ptr = nullptr;
|
||||
if (tid2eid.has_value()) {
|
||||
input_ids_ptr = input_ids.value().data_ptr<int64_t>();
|
||||
tid2eid_ptr = tid2eid.value().data_ptr<int64_t>();
|
||||
}
|
||||
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
|
||||
gating_output, topk_weights.data_ptr<float>(),
|
||||
topk_indices.data_ptr<int64_t>(), token_expert_indices.data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, use_hash, input_ids_ptr, tid2eid_ptr, stream);
|
||||
}
|
||||
}
|
||||
|
||||
void topk_softplus_sqrt(
|
||||
torch::Tensor& topk_weights, // [num_tokens, topk]
|
||||
torch::Tensor& topk_indices, // [num_tokens, topk]
|
||||
torch::Tensor& token_expert_indices, // [num_tokens, topk]
|
||||
torch::Tensor& gating_output, // [num_tokens, num_experts]
|
||||
bool renormalize, double routed_scaling_factor,
|
||||
const c10::optional<torch::Tensor>& correction_bias,
|
||||
const c10::optional<torch::Tensor>& input_ids,
|
||||
const c10::optional<torch::Tensor>& tid2eid) {
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
const int topk = topk_weights.size(-1);
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
if (gating_output.scalar_type() == at::ScalarType::Float) {
|
||||
dispatch_topk_softplus_sqrt_launch<float>(
|
||||
gating_output.data_ptr<float>(), topk_weights, topk_indices,
|
||||
token_expert_indices, num_tokens, num_experts, topk, renormalize,
|
||||
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
|
||||
} else if (gating_output.scalar_type() == at::ScalarType::Half) {
|
||||
dispatch_topk_softplus_sqrt_launch<__half>(
|
||||
reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
|
||||
topk_weights, topk_indices, token_expert_indices, num_tokens,
|
||||
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid, stream);
|
||||
} else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
|
||||
dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(
|
||||
gating_output.data_ptr<at::BFloat16>()),
|
||||
topk_weights, topk_indices, token_expert_indices, num_tokens,
|
||||
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid, stream);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unsupported gating_output data type: ",
|
||||
gating_output.scalar_type());
|
||||
}
|
||||
}
|
||||
@@ -16,6 +16,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
|
||||
"bias) -> ()");
|
||||
m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
m.def(
|
||||
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, float "
|
||||
"routed_scaling_factor, Tensor? "
|
||||
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
|
||||
m.impl("topk_softplus_sqrt", torch::kCUDA, &topk_softplus_sqrt);
|
||||
#endif
|
||||
// Calculate the result of moe by summing up the partial results
|
||||
// from all selected experts.
|
||||
m.def("moe_sum(Tensor input, Tensor! output) -> ()");
|
||||
|
||||
+9
-1
@@ -100,6 +100,11 @@ void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
|
||||
bool is_neox, torch::Tensor& position_ids,
|
||||
int64_t forced_token_heads_per_warp);
|
||||
|
||||
void fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
|
||||
torch::Tensor& q, torch::Tensor const& kv, torch::Tensor& k_cache,
|
||||
torch::Tensor const& slot_mapping, torch::Tensor const& position_ids,
|
||||
torch::Tensor const& cos_sin_cache, double eps, int64_t cache_block_size);
|
||||
|
||||
void apply_repetition_penalties_(torch::Tensor& logits,
|
||||
const torch::Tensor& prompt_mask,
|
||||
const torch::Tensor& output_mask,
|
||||
@@ -153,10 +158,13 @@ void silu_and_mul_per_block_quant(torch::Tensor& out,
|
||||
|
||||
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
|
||||
std::optional<torch::Tensor> key, int64_t head_size,
|
||||
torch::Tensor& cos_sin_cache, bool is_neox);
|
||||
torch::Tensor& cos_sin_cache, bool is_neox,
|
||||
int64_t rope_dim_offset, bool inverse);
|
||||
|
||||
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
|
||||
|
||||
void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
|
||||
|
||||
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
|
||||
torch::Tensor& scale);
|
||||
|
||||
|
||||
+17
-16
@@ -18,7 +18,6 @@ namespace persistent {
|
||||
// Constants
|
||||
// ============================================================================
|
||||
|
||||
constexpr int TopK = 2048;
|
||||
constexpr int kThreadsPerBlock = 1024;
|
||||
constexpr int RADIX = 256;
|
||||
|
||||
@@ -128,11 +127,12 @@ struct RadixRowState {
|
||||
|
||||
struct PersistentTopKParams {
|
||||
const float* __restrict__ input; // [num_rows, stride]
|
||||
int32_t* __restrict__ output; // [num_rows, TopK]
|
||||
int32_t* __restrict__ output; // [num_rows, top_k]
|
||||
int32_t* __restrict__ lengths; // [num_rows]
|
||||
RadixRowState* row_states; // large path: per-group state
|
||||
uint32_t num_rows;
|
||||
uint32_t stride;
|
||||
uint32_t top_k; // actual k value for output stride
|
||||
uint32_t chunk_size; // large path: elements per CTA
|
||||
uint32_t ctas_per_group; // 1=medium, >1=large
|
||||
uint32_t max_seq_len; // max seq_len across all rows (for early CTA exit)
|
||||
@@ -154,6 +154,7 @@ __device__ __forceinline__ uint32_t decode_bin(float x) {
|
||||
return key >> 5;
|
||||
}
|
||||
|
||||
template <int TopK>
|
||||
__device__ __noinline__ void histogram_2048_topk(
|
||||
const float* __restrict__ logits, int32_t* __restrict__ output_indices,
|
||||
int32_t seq_len) {
|
||||
@@ -418,6 +419,7 @@ __device__ __noinline__ void histogram_2048_topk(
|
||||
// by: DarkSharpness
|
||||
// which at the same time is an optimized topk kernel copied from tilelang
|
||||
// kernel
|
||||
template <int TopK>
|
||||
__device__ __noinline__ void histogram_256_topk(
|
||||
const float* __restrict__ logits, int* __restrict__ output_indices,
|
||||
int logits_offset, int seq_len) {
|
||||
@@ -649,7 +651,7 @@ __device__ __forceinline__ void wait_ge(int* ptr, int target_val,
|
||||
// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215
|
||||
// ============================================================================
|
||||
|
||||
template <uint32_t VEC_SIZE>
|
||||
template <int TopK, uint32_t VEC_SIZE>
|
||||
__device__ void radix_topk(const float* __restrict__ row_input,
|
||||
int32_t* __restrict__ row_output, uint32_t seq_len,
|
||||
uint32_t my_chunk_start, uint32_t chunk_size,
|
||||
@@ -857,7 +859,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
|
||||
// see filtered_topk.cuh)
|
||||
// ============================================================================
|
||||
|
||||
template <uint32_t VEC_SIZE = 1>
|
||||
template <int TopK = 2048, uint32_t VEC_SIZE = 1>
|
||||
__global__ void __launch_bounds__(kThreadsPerBlock, 2)
|
||||
persistent_topk_kernel(PersistentTopKParams params) {
|
||||
const uint32_t tx = threadIdx.x;
|
||||
@@ -915,7 +917,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
|
||||
if (row_idx >= params.num_rows) break;
|
||||
|
||||
const uint32_t seq_len = params.lengths[row_idx];
|
||||
int32_t* row_output = params.output + row_idx * TopK;
|
||||
int32_t* row_output = params.output + row_idx * params.top_k;
|
||||
const float* row_input = params.input + row_idx * params.stride;
|
||||
|
||||
if (seq_len <= RADIX_THRESHOLD) {
|
||||
@@ -927,19 +929,19 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
|
||||
row_output[i] = (i < seq_len) ? static_cast<int32_t>(i) : -1;
|
||||
}
|
||||
} else if (seq_len <= static_cast<uint32_t>(HIST2048_THRESHOLD)) {
|
||||
histogram_2048_topk(row_input, row_output, seq_len);
|
||||
histogram_2048_topk<TopK>(row_input, row_output, seq_len);
|
||||
} else {
|
||||
histogram_256_topk(row_input, row_output, 0, seq_len);
|
||||
histogram_256_topk<TopK>(row_input, row_output, 0, seq_len);
|
||||
}
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
const uint32_t my_chunk_start = cta_in_group * chunk_size;
|
||||
radix_topk<VEC_SIZE>(row_input, row_output, seq_len, my_chunk_start,
|
||||
chunk_size, local_histogram, suffix_sum,
|
||||
shared_scalars, shared_ordered, state, cta_in_group,
|
||||
ctas_per_group, barrier_phase, iter, tx);
|
||||
radix_topk<TopK, VEC_SIZE>(
|
||||
row_input, row_output, seq_len, my_chunk_start, chunk_size,
|
||||
local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
|
||||
cta_in_group, ctas_per_group, barrier_phase, iter, tx);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1011,7 +1013,6 @@ struct FilteredTopKTraits<float> {
|
||||
}
|
||||
};
|
||||
|
||||
constexpr uint32_t FILTERED_TOPK_MAX_K = 2048;
|
||||
constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024;
|
||||
constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE =
|
||||
16 * 1024; // 16K indices per buffer
|
||||
@@ -1025,7 +1026,7 @@ constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC =
|
||||
* \tparam IdType Index type (int32_t)
|
||||
* \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8)
|
||||
*/
|
||||
template <typename DType, typename IdType, int VEC_SIZE>
|
||||
template <typename DType, typename IdType, int VEC_SIZE, uint32_t MAX_K = 2048>
|
||||
__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
|
||||
FilteredTopKUnifiedKernel(const DType* __restrict__ input,
|
||||
IdType* __restrict__ output,
|
||||
@@ -1059,7 +1060,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS)
|
||||
alignas(128) __shared__ int s_counter;
|
||||
alignas(128) __shared__ int s_threshold_bin_id;
|
||||
alignas(128) __shared__ int s_num_input[2];
|
||||
alignas(128) __shared__ int s_indices[FILTERED_TOPK_MAX_K];
|
||||
alignas(128) __shared__ int s_indices[MAX_K];
|
||||
|
||||
auto& s_histogram = s_histogram_buf[0];
|
||||
|
||||
@@ -1280,7 +1281,7 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) {
|
||||
return static_cast<int>(g);
|
||||
}
|
||||
|
||||
template <typename DType, typename IdType>
|
||||
template <typename DType, typename IdType, uint32_t MAX_K = 2048>
|
||||
cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
|
||||
IdType* lengths, uint32_t num_rows,
|
||||
uint32_t top_k_val, uint32_t max_len,
|
||||
@@ -1297,7 +1298,7 @@ cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices,
|
||||
|
||||
#define DISPATCH_VEC_SIZE(VS) \
|
||||
if (vec_size == VS) { \
|
||||
auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS>; \
|
||||
auto kernel = FilteredTopKUnifiedKernel<DType, IdType, VS, MAX_K>; \
|
||||
FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \
|
||||
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \
|
||||
FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \
|
||||
|
||||
@@ -9,28 +9,29 @@ namespace vllm {
|
||||
|
||||
template <typename scalar_t, bool IS_NEOX>
|
||||
inline __device__ void apply_token_rotary_embedding(
|
||||
scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr,
|
||||
const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim) {
|
||||
scalar_t* __restrict__ arr, const float* __restrict__ cos_ptr,
|
||||
const float* __restrict__ sin_ptr, int rot_offset, int embed_dim,
|
||||
const bool inverse) {
|
||||
int x_index, y_index;
|
||||
scalar_t cos, sin;
|
||||
float cos_f, sin_f;
|
||||
if (IS_NEOX) {
|
||||
// GPT-NeoX style rotary embedding.
|
||||
x_index = rot_offset;
|
||||
y_index = embed_dim + rot_offset;
|
||||
cos = VLLM_LDG(cos_ptr + x_index);
|
||||
sin = VLLM_LDG(sin_ptr + x_index);
|
||||
cos_f = VLLM_LDG(cos_ptr + x_index);
|
||||
sin_f = VLLM_LDG(sin_ptr + x_index);
|
||||
} else {
|
||||
// GPT-J style rotary embedding.
|
||||
x_index = 2 * rot_offset;
|
||||
y_index = 2 * rot_offset + 1;
|
||||
cos = VLLM_LDG(cos_ptr + x_index / 2);
|
||||
sin = VLLM_LDG(sin_ptr + x_index / 2);
|
||||
cos_f = VLLM_LDG(cos_ptr + x_index / 2);
|
||||
sin_f = VLLM_LDG(sin_ptr + x_index / 2);
|
||||
}
|
||||
|
||||
const scalar_t x = arr[x_index];
|
||||
const scalar_t y = arr[y_index];
|
||||
arr[x_index] = x * cos - y * sin;
|
||||
arr[y_index] = y * cos + x * sin;
|
||||
if (inverse) {
|
||||
sin_f = -sin_f;
|
||||
}
|
||||
const float x_f = static_cast<float>(arr[x_index]);
|
||||
const float y_f = static_cast<float>(arr[y_index]);
|
||||
arr[x_index] = static_cast<scalar_t>(x_f * cos_f - y_f * sin_f);
|
||||
arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f);
|
||||
}
|
||||
|
||||
template <typename scalar_t, bool IS_NEOX>
|
||||
@@ -42,22 +43,23 @@ inline __device__ void apply_rotary_embedding(
|
||||
// [batch_size, seq_len, num_kv_heads,
|
||||
// head_size] or [num_tokens, num_kv_heads,
|
||||
// head_size]
|
||||
const scalar_t* cache_ptr, const int head_size, const int num_heads,
|
||||
const float* 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 head_stride, const int64_t rope_dim_offset,
|
||||
const bool inverse) {
|
||||
const int embed_dim = rot_dim / 2;
|
||||
const scalar_t* cos_ptr = cache_ptr;
|
||||
const scalar_t* sin_ptr = cache_ptr + embed_dim;
|
||||
const float* cos_ptr = cache_ptr;
|
||||
const float* sin_ptr = cache_ptr + embed_dim;
|
||||
|
||||
const int nq = num_heads * embed_dim;
|
||||
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
|
||||
const int head_idx = i / embed_dim;
|
||||
const int64_t token_head =
|
||||
token_idx * query_stride + head_idx * head_stride;
|
||||
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>(
|
||||
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
|
||||
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
|
||||
}
|
||||
|
||||
if (key != nullptr) {
|
||||
@@ -65,10 +67,10 @@ inline __device__ void apply_rotary_embedding(
|
||||
for (int i = threadIdx.x; i < nk; i += blockDim.x) {
|
||||
const int head_idx = i / embed_dim;
|
||||
const int64_t token_head =
|
||||
token_idx * key_stride + head_idx * head_stride;
|
||||
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>(
|
||||
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
|
||||
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -84,19 +86,18 @@ __global__ void rotary_embedding_kernel(
|
||||
// [batch_size, seq_len, num_kv_heads,
|
||||
// head_size] or [num_tokens, num_kv_heads,
|
||||
// head_size]
|
||||
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
|
||||
// 2]
|
||||
const float* __restrict__ cos_sin_cache, // [max_position, rot_dim] fp32
|
||||
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) {
|
||||
// Each thread block is responsible for one token.
|
||||
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 scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
|
||||
const float* cache_ptr = cos_sin_cache + pos * rot_dim;
|
||||
|
||||
apply_rotary_embedding<scalar_t, IS_NEOX>(
|
||||
query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
|
||||
token_idx, query_stride, key_stride, head_stride);
|
||||
token_idx, query_stride, key_stride, head_stride, rope_dim_offset,
|
||||
inverse);
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
@@ -115,7 +116,7 @@ void rotary_embedding(
|
||||
// [num_tokens, num_heads, head_size]
|
||||
int64_t head_size,
|
||||
torch::Tensor& cos_sin_cache, // [max_position, rot_dim]
|
||||
bool is_neox) {
|
||||
bool is_neox, int64_t rope_dim_offset, bool inverse) {
|
||||
// num_tokens = batch_size * seq_len
|
||||
int64_t num_tokens = positions.numel();
|
||||
int positions_ndim = positions.dim();
|
||||
@@ -154,6 +155,8 @@ void rotary_embedding(
|
||||
int seq_dim_idx = positions_ndim - 1;
|
||||
int64_t query_stride = query.stride(seq_dim_idx);
|
||||
int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
|
||||
|
||||
TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size);
|
||||
// Determine head stride: for [*, heads, head_size] use stride of last dim;
|
||||
// for flat [*, heads*head_size], heads blocks are contiguous of size
|
||||
// head_size
|
||||
@@ -165,20 +168,23 @@ 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,
|
||||
cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride, key_stride,
|
||||
head_stride, num_heads, num_kv_heads, head_size);
|
||||
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,
|
||||
cos_sin_cache.data_ptr<scalar_t>(), rot_dim, query_stride,
|
||||
key_stride, head_stride, num_heads, num_kv_heads, head_size);
|
||||
cache_f32.data_ptr<float>(), rot_dim, query_stride, key_stride,
|
||||
head_stride, num_heads, num_kv_heads, head_size, rope_dim_offset,
|
||||
inverse);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
+2
-11
@@ -40,15 +40,6 @@ using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz;
|
||||
#define __HIP__FP8MFMA__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx1100__) || defined(__gfx1101__) || \
|
||||
defined(__gfx1150__) || defined(__gfx1151__))
|
||||
#define __HIP__GFX11__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
|
||||
#define __HIP__GFX12__
|
||||
#endif
|
||||
|
||||
#if defined(NDEBUG)
|
||||
#undef NDEBUG
|
||||
#include <assert.h>
|
||||
@@ -1629,7 +1620,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
|
||||
}
|
||||
}
|
||||
|
||||
#elif defined(__HIP__GFX11__)
|
||||
#elif defined(__GFX11__)
|
||||
|
||||
using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
|
||||
|
||||
@@ -2388,7 +2379,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
|
||||
out_ptr[threadIdx.x] = from_float<scalar_t>(acc);
|
||||
}
|
||||
|
||||
#elif defined(__HIP__GFX12__)
|
||||
#elif defined(__GFX12__)
|
||||
|
||||
using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float;
|
||||
|
||||
|
||||
+18
-23
@@ -26,16 +26,11 @@
|
||||
#define __HIP__GFX9__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && \
|
||||
(defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1150__) || \
|
||||
defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
|
||||
// Combined RDNA macro (gfx11 + gfx12) - both use 32-wide wavefronts
|
||||
#if defined(__GFX11__) || defined(__GFX12__)
|
||||
#define __HIP__GFX1X__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__))
|
||||
#define __HIP__GFX12__
|
||||
#endif
|
||||
|
||||
#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__))
|
||||
#define __HIP__MI3XX__
|
||||
#endif
|
||||
@@ -1845,7 +1840,7 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
|
||||
return out_c;
|
||||
}
|
||||
|
||||
#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
#if defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
@@ -1893,7 +1888,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
float sB = *s_B;
|
||||
|
||||
while (m < M) {
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
|
||||
float sum[N][YTILE] = {};
|
||||
#else
|
||||
@@ -1931,7 +1926,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#pragma unroll
|
||||
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
|
||||
for (uint32_t n = 0; n < N; n++) {
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
|
||||
for (int y = 0; y < YTILE; ++y) {
|
||||
#pragma unroll
|
||||
@@ -1955,7 +1950,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
}
|
||||
|
||||
// Final reduction
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
@@ -1993,7 +1988,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#endif
|
||||
|
||||
const bool writeback_lane =
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
threadIdx.x == (THRDS - 1);
|
||||
#else
|
||||
threadIdx.x == 0;
|
||||
@@ -2009,7 +2004,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
if (y + m >= M) break; // To avoid mem access fault.
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
float result = sum[n][y] * sA * sB;
|
||||
#else
|
||||
float result = sum[n][y][0] * sA * sB;
|
||||
@@ -2027,7 +2022,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
m += CuCount * _WvPrGrp * YTILE;
|
||||
}
|
||||
}
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
|
||||
@@ -2039,9 +2034,9 @@ __global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp,
|
||||
const int _WvPrGrp, const int CuCount) {
|
||||
UNREACHABLE_CODE
|
||||
}
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
|
||||
#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
#if defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
@@ -2088,7 +2083,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
float sB = *s_B;
|
||||
|
||||
while (m < M) {
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8
|
||||
float sum[N][YTILE] = {};
|
||||
#else
|
||||
@@ -2128,7 +2123,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#pragma unroll
|
||||
for (uint32_t k2 = 0; k2 < UNRL; k2++) {
|
||||
for (uint32_t n = 0; n < N; n++) {
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4)
|
||||
for (int y = 0; y < YTILE; ++y) {
|
||||
#pragma unroll
|
||||
@@ -2152,7 +2147,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
}
|
||||
|
||||
// Final reduction
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
// gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
@@ -2190,7 +2185,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
#endif
|
||||
|
||||
const bool writeback_lane =
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
threadIdx.x == (THRDS - 1);
|
||||
#else
|
||||
threadIdx.x == 0;
|
||||
@@ -2206,7 +2201,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
for (int n = 0; n < N; n++) {
|
||||
for (int y = 0; y < YTILE; y++) {
|
||||
if (y + m >= M) break; // To avoid mem access fault.
|
||||
#ifdef __HIP__GFX12__
|
||||
#ifdef __GFX12__
|
||||
float result = sum[n][y] * sA * sB;
|
||||
#else
|
||||
float result = sum[n][y][0] * sA * sB;
|
||||
@@ -2224,7 +2219,7 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
|
||||
m += CuCount * _WvPrGrp * YTILE;
|
||||
}
|
||||
}
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__)
|
||||
#else // !defined(__HIP__MI3XX__) && !defined(__GFX12__)
|
||||
template <typename scalar_t, typename fp8_t, int THRDS, int YTILE, int WvPrGrp,
|
||||
int A_CHUNK, int UNRL, int N>
|
||||
__global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
|
||||
@@ -2236,7 +2231,7 @@ __global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp,
|
||||
const int CuCount) {
|
||||
UNREACHABLE_CODE
|
||||
}
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__)
|
||||
#endif // defined(__HIP__MI3XX__) || defined(__GFX12__)
|
||||
|
||||
void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a,
|
||||
const std::optional<at::Tensor>& in_bias, at::Tensor& out_c,
|
||||
|
||||
+7
-1
@@ -258,7 +258,13 @@ __device__ bool processHistogramStep(
|
||||
auto processBins = [&](float logit, int idx) {
|
||||
if (isPartialMatch<patternShift>(logit, logitPattern)) {
|
||||
uint32_t binIdx = extractBinIdx<step>(logit);
|
||||
if (binIdx < thresholdBinIdx) {
|
||||
// Only write elements with binIdx < thresholdBinIdx when:
|
||||
// 1. This is step 0 and the threshold bin is small enough (no step 1)
|
||||
// 2. This is step >= 1 (where pattern matching filters correctly)
|
||||
// This prevents duplicates when step 0 and step 1 both run.
|
||||
bool shouldWriteDirectly =
|
||||
(step == 0 && smemFinalBinSize[0] <= kNumFinalItems) || (step >= 1);
|
||||
if (binIdx < thresholdBinIdx && shouldWriteDirectly) {
|
||||
// The element is part of the top-k selection
|
||||
int dstIdx = atomicAdd(&smemFoundTopKValues[0], 1);
|
||||
|
||||
|
||||
+59
-35
@@ -10,33 +10,17 @@
|
||||
#include "persistent_topk.cuh"
|
||||
#endif
|
||||
|
||||
void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
torch::Tensor& output, torch::Tensor& workspace, int64_t k,
|
||||
int64_t max_seq_len) {
|
||||
namespace {
|
||||
|
||||
#ifndef USE_ROCM
|
||||
TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
|
||||
TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
|
||||
TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
|
||||
TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
|
||||
TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
|
||||
TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
|
||||
TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
|
||||
TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
|
||||
"lengths must be 1D or 2D");
|
||||
TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
|
||||
TORCH_CHECK(output.dim() == 2, "output must be 2D");
|
||||
template <int TopK>
|
||||
void launch_persistent_topk(const torch::Tensor& logits,
|
||||
const torch::Tensor& lengths, torch::Tensor& output,
|
||||
torch::Tensor& workspace, int64_t max_seq_len) {
|
||||
namespace P = vllm::persistent;
|
||||
|
||||
const int64_t num_rows = logits.size(0);
|
||||
const int64_t stride = logits.size(1);
|
||||
|
||||
TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
|
||||
TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
|
||||
"output size mismatch");
|
||||
namespace P = vllm::persistent;
|
||||
|
||||
TORCH_CHECK(k == P::TopK, "k must be 2048");
|
||||
TORCH_CHECK(k <= stride, "k out of range");
|
||||
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
static int num_sms = 0;
|
||||
@@ -50,18 +34,17 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
}
|
||||
|
||||
if (num_rows > 32 && max_smem_per_block >= 128 * 1024) {
|
||||
cudaError_t status = vllm::FilteredTopKRaggedTransform<float, int32_t>(
|
||||
logits.data_ptr<float>(), output.data_ptr<int32_t>(),
|
||||
lengths.data_ptr<int32_t>(), static_cast<uint32_t>(num_rows),
|
||||
static_cast<uint32_t>(k), static_cast<uint32_t>(stride), stream);
|
||||
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 failed: ", cudaGetErrorString(status));
|
||||
} else {
|
||||
TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
|
||||
TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8");
|
||||
|
||||
// Smem cap: smaller smem → more CTAs/group → more per-row parallelism for
|
||||
// large path. Empirically tuned.
|
||||
int effective_max_smem;
|
||||
if (num_rows <= 4) {
|
||||
effective_max_smem =
|
||||
@@ -101,7 +84,7 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
|
||||
int occupancy = 1;
|
||||
cudaOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
&occupancy, P::persistent_topk_kernel<4>, P::kThreadsPerBlock,
|
||||
&occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock,
|
||||
smem_size);
|
||||
if (occupancy < 1) occupancy = 1;
|
||||
|
||||
@@ -121,15 +104,16 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
params.lengths = lengths.data_ptr<int32_t>();
|
||||
params.num_rows = static_cast<uint32_t>(num_rows);
|
||||
params.stride = static_cast<uint32_t>(stride);
|
||||
params.top_k = static_cast<uint32_t>(TopK);
|
||||
params.chunk_size = chunk_size;
|
||||
params.row_states =
|
||||
reinterpret_cast<P::RadixRowState*>(workspace.data_ptr<uint8_t>());
|
||||
params.ctas_per_group = ctas_per_group;
|
||||
params.max_seq_len = static_cast<uint32_t>(max_seq_len);
|
||||
|
||||
#define LAUNCH_PERSISTENT(VS) \
|
||||
#define LAUNCH_PERSISTENT(TOPK_VAL, VS) \
|
||||
do { \
|
||||
auto kernel = &P::persistent_topk_kernel<VS>; \
|
||||
auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>; \
|
||||
cudaError_t err = cudaFuncSetAttribute( \
|
||||
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \
|
||||
TORCH_CHECK(err == cudaSuccess, \
|
||||
@@ -138,11 +122,11 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
} while (0)
|
||||
|
||||
if (vec_size == 4) {
|
||||
LAUNCH_PERSISTENT(4);
|
||||
LAUNCH_PERSISTENT(TopK, 4);
|
||||
} else if (vec_size == 2) {
|
||||
LAUNCH_PERSISTENT(2);
|
||||
LAUNCH_PERSISTENT(TopK, 2);
|
||||
} else {
|
||||
LAUNCH_PERSISTENT(1);
|
||||
LAUNCH_PERSISTENT(TopK, 1);
|
||||
}
|
||||
#undef LAUNCH_PERSISTENT
|
||||
}
|
||||
@@ -150,6 +134,46 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
cudaError_t err = cudaGetLastError();
|
||||
TORCH_CHECK(err == cudaSuccess,
|
||||
"persistent_topk failed: ", cudaGetErrorString(err));
|
||||
}
|
||||
#endif
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths,
|
||||
torch::Tensor& output, torch::Tensor& workspace, int64_t k,
|
||||
int64_t max_seq_len) {
|
||||
#ifndef USE_ROCM
|
||||
TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
|
||||
TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
|
||||
TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
|
||||
TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported");
|
||||
TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32");
|
||||
TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32");
|
||||
TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
|
||||
TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
|
||||
"lengths must be 1D or 2D");
|
||||
TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
|
||||
TORCH_CHECK(output.dim() == 2, "output must be 2D");
|
||||
|
||||
const int64_t num_rows = logits.size(0);
|
||||
const int64_t stride = logits.size(1);
|
||||
|
||||
TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
|
||||
TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
|
||||
"output size mismatch");
|
||||
TORCH_CHECK(k == 512 || k == 1024 || k == 2048,
|
||||
"persistent_topk supports k=512, k=1024, or k=2048, got k=", k);
|
||||
|
||||
if (k == 512) {
|
||||
launch_persistent_topk<512>(logits, lengths, output, workspace,
|
||||
max_seq_len);
|
||||
} else if (k == 1024) {
|
||||
launch_persistent_topk<1024>(logits, lengths, output, workspace,
|
||||
max_seq_len);
|
||||
} else {
|
||||
launch_persistent_topk<2048>(logits, lengths, output, workspace,
|
||||
max_seq_len);
|
||||
}
|
||||
#else
|
||||
TORCH_CHECK(false, "persistent_topk is not supported on ROCm");
|
||||
#endif
|
||||
|
||||
+21
-1
@@ -106,6 +106,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
|
||||
ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
|
||||
|
||||
// SwiGLU activation with input clamping.
|
||||
ops.def(
|
||||
"silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
|
||||
"-> ()");
|
||||
ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
|
||||
|
||||
ops.def(
|
||||
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
|
||||
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
|
||||
@@ -177,6 +183,19 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"int forced_token_heads_per_warp=-1) -> ()");
|
||||
ops.impl("fused_qk_norm_rope", torch::kCUDA, &fused_qk_norm_rope);
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// Horizontally-fused DeepseekV4-MLA: per-head RMSNorm + GPT-J RoPE for Q, and
|
||||
// GPT-J RoPE + UE8M0 FP8 quant + paged cache insert for KV, all in one
|
||||
// kernel launch.
|
||||
ops.def(
|
||||
"fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert("
|
||||
"Tensor! q, Tensor kv, Tensor! k_cache, "
|
||||
"Tensor slot_mapping, Tensor position_ids, Tensor cos_sin_cache, "
|
||||
"float eps, int cache_block_size) -> ()");
|
||||
ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA,
|
||||
&fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert);
|
||||
#endif
|
||||
|
||||
// Apply repetition penalties to logits in-place
|
||||
ops.def(
|
||||
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
|
||||
@@ -240,7 +259,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def(
|
||||
"rotary_embedding(Tensor positions, Tensor! query,"
|
||||
" Tensor!? key, int head_size,"
|
||||
" Tensor cos_sin_cache, bool is_neox) -> ()");
|
||||
" Tensor cos_sin_cache, bool is_neox, int "
|
||||
"rope_dim_offset=0, bool inverse=False) -> ()");
|
||||
ops.impl("rotary_embedding", torch::kCUDA, &rotary_embedding);
|
||||
|
||||
// Quantization ops
|
||||
|
||||
+15
-8
@@ -22,7 +22,7 @@
|
||||
# docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
|
||||
# =============================================================================
|
||||
|
||||
ARG CUDA_VERSION=13.0.0
|
||||
ARG CUDA_VERSION=13.0.2
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG UBUNTU_VERSION=22.04
|
||||
|
||||
@@ -188,7 +188,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# Explicitly set the list to avoid issues with torch 2.2
|
||||
# See https://github.com/pytorch/pytorch/pull/123243
|
||||
# From versions.json: .torch.cuda_arch_list
|
||||
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
|
||||
ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
|
||||
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
#################### BUILD BASE IMAGE ####################
|
||||
|
||||
@@ -538,9 +538,11 @@ RUN CUDA_VERSION_DASH=$(echo $CUDA_VERSION | cut -d. -f1,2 | tr '.' '-') && \
|
||||
cuda-nvrtc-${CUDA_VERSION_DASH} \
|
||||
cuda-cuobjdump-${CUDA_VERSION_DASH} \
|
||||
libcurand-dev-${CUDA_VERSION_DASH} \
|
||||
libcublas-${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).
|
||||
@@ -765,7 +772,7 @@ ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# install kv_connectors if requested
|
||||
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.9 9.0 10.0 12.0'
|
||||
ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX'
|
||||
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,source=requirements/kv_connectors.txt,target=/tmp/kv_connectors.txt,ro \
|
||||
|
||||
@@ -77,7 +77,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system $pkgs --index-url https://download.pytorch.org/whl/nightly/cu128
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system numba==0.61.2
|
||||
uv pip install --system numba==0.65.0
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system -r requirements/common.txt
|
||||
|
||||
+38
-5
@@ -124,10 +124,10 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
|
||||
|
||||
# RIXL/UCX build stages
|
||||
FROM base AS build_rixl
|
||||
ARG RIXL_BRANCH="f33a5599"
|
||||
ARG RIXL_BRANCH="bf4a7214"
|
||||
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
|
||||
ARG UCX_BRANCH="da3fac2a"
|
||||
ARG UCX_REPO="https://github.com/ROCm/ucx.git"
|
||||
ARG UCX_BRANCH="7009d7a1"
|
||||
ARG UCX_REPO="https://github.com/openucx/ucx.git"
|
||||
ENV ROCM_PATH=/opt/rocm
|
||||
ENV UCX_HOME=/usr/local/ucx
|
||||
ENV RIXL_HOME=/usr/local/rixl
|
||||
@@ -165,7 +165,7 @@ RUN cd /usr/local/src && \
|
||||
--disable-doxygen-doc \
|
||||
--enable-optimizations \
|
||||
--enable-devel-headers \
|
||||
--with-rocm=/opt/rocm \
|
||||
--with-rocm=${ROCM_PATH} \
|
||||
--with-verbs \
|
||||
--with-dm \
|
||||
--enable-mt && \
|
||||
@@ -186,7 +186,12 @@ RUN git clone ${RIXL_REPO} /opt/rixl && \
|
||||
ninja install
|
||||
|
||||
# Generate RIXL wheel
|
||||
RUN cd /opt/rixl && mkdir -p /app/install && \
|
||||
# Exclude libcore and libpull from auditwheel: transitive dependencies
|
||||
# that are not shipped in the wheel and vary across base images.
|
||||
RUN cd /opt/rixl && \
|
||||
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
|
||||
contrib/build-wheel.sh && \
|
||||
mkdir -p /app/install && \
|
||||
./contrib/build-wheel.sh \
|
||||
--output-dir /app/install \
|
||||
--rocm-dir ${ROCM_PATH} \
|
||||
@@ -307,6 +312,30 @@ RUN --mount=type=bind,source=.git,target=vllm/.git \
|
||||
&& echo "Detected vLLM version: ${VLLM_VERSION}" \
|
||||
&& echo "${VLLM_VERSION}" > /tmp/vllm_version.txt
|
||||
|
||||
# Fail if git-based package dependencies are found in requirements files
|
||||
# (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI)
|
||||
# Extra notes: pip install is able to handle git+ URLs, but uv doesn't.
|
||||
RUN echo "Checking for git-based packages in requirements files..." \
|
||||
&& echo "Checking common.txt for git-based packages:" \
|
||||
&& if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; then \
|
||||
echo "ERROR: Git-based packages found in common.txt:"; \
|
||||
grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/common.txt; \
|
||||
echo "Please publish these packages to PyPI instead of using git dependencies."; \
|
||||
exit 1; \
|
||||
else \
|
||||
echo " ✓ No git-based packages found in common.txt"; \
|
||||
fi \
|
||||
&& echo "Checking rocm.txt for git-based packages:" \
|
||||
&& if grep -q 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; then \
|
||||
echo "ERROR: Git-based packages found in rocm.txt:"; \
|
||||
grep 'git+' ${COMMON_WORKDIR}/vllm/requirements/rocm.txt; \
|
||||
echo "Please publish these packages to PyPI instead of using git dependencies."; \
|
||||
exit 1; \
|
||||
else \
|
||||
echo " ✓ No git-based packages found in rocm.txt"; \
|
||||
fi \
|
||||
&& echo "All requirements files are clean - no git-based packages found"
|
||||
|
||||
# Pin vLLM dependencies to exact versions of custom ROCm wheels
|
||||
# This ensures 'pip install vllm' automatically installs correct torch/triton/torchvision/amdsmi
|
||||
COPY tools/vllm-rocm/pin_rocm_dependencies.py /tmp/pin_rocm_dependencies.py
|
||||
@@ -407,6 +436,10 @@ COPY --from=export_vllm /vllm_v1 /usr/local/lib/python${PYTHON_VERSION}/dist-pac
|
||||
ENV MIOPEN_DEBUG_CONV_DIRECT=0
|
||||
ENV MIOPEN_DEBUG_CONV_GEMM=0
|
||||
|
||||
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
|
||||
# See: https://github.com/ROCm/rocm-libraries/issues/6266
|
||||
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
|
||||
|
||||
# Source code is used in the `python_only_compile.sh` test
|
||||
# We hide it inside `src/` so that this source code
|
||||
# will not be imported by other tests
|
||||
|
||||
@@ -50,9 +50,9 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
|
||||
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
|
||||
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
|
||||
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
|
||||
rm "${ONECCL_INSTALLER}" && \
|
||||
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
|
||||
@@ -164,7 +164,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# FIX triton
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.6.0
|
||||
uv pip install triton-xpu==3.7.0
|
||||
|
||||
# remove torch bundled oneccl to avoid conflicts
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
@@ -20,7 +20,7 @@ variable "NVCC_THREADS" {
|
||||
}
|
||||
|
||||
variable "TORCH_CUDA_ARCH_LIST" {
|
||||
default = "8.0 8.9 9.0 10.0"
|
||||
default = "8.0 8.9 9.0 10.0 11.0 12.0"
|
||||
}
|
||||
|
||||
variable "COMMIT" {
|
||||
@@ -88,7 +88,6 @@ target "test-ubuntu2404" {
|
||||
args = {
|
||||
UBUNTU_VERSION = "24.04"
|
||||
GDRCOPY_OS_VERSION = "Ubuntu24_04"
|
||||
FLASHINFER_AOT_COMPILE = "true"
|
||||
}
|
||||
output = ["type=docker"]
|
||||
}
|
||||
@@ -100,7 +99,6 @@ target "openai-ubuntu2404" {
|
||||
args = {
|
||||
UBUNTU_VERSION = "24.04"
|
||||
GDRCOPY_OS_VERSION = "Ubuntu24_04"
|
||||
FLASHINFER_AOT_COMPILE = "true"
|
||||
}
|
||||
output = ["type=docker"]
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
"_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
|
||||
"variable": {
|
||||
"CUDA_VERSION": {
|
||||
"default": "13.0.0"
|
||||
"default": "13.0.2"
|
||||
},
|
||||
"PYTHON_VERSION": {
|
||||
"default": "3.12"
|
||||
@@ -11,10 +11,10 @@
|
||||
"default": "22.04"
|
||||
},
|
||||
"BUILD_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.0-devel-ubuntu22.04"
|
||||
"default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
|
||||
},
|
||||
"FINAL_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.0-base-ubuntu22.04"
|
||||
"default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
|
||||
},
|
||||
"GET_PIP_URL": {
|
||||
"default": "https://bootstrap.pypa.io/get-pip.py"
|
||||
@@ -32,7 +32,7 @@
|
||||
"default": "false"
|
||||
},
|
||||
"TORCH_CUDA_ARCH_LIST": {
|
||||
"default": "7.0 7.5 8.0 8.9 9.0 10.0 12.0"
|
||||
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0+PTX"
|
||||
},
|
||||
"MAX_JOBS": {
|
||||
"default": "2"
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 315 KiB After Width: | Height: | Size: 315 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 108 KiB |
File diff suppressed because one or more lines are too long
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 |
@@ -108,6 +108,38 @@ P99 ITL (ms): 8.39
|
||||
==================================================
|
||||
```
|
||||
|
||||
#### Results Visualization
|
||||
|
||||
The `--plot-timeline` and `--plot-dataset-stats` can be used to generate respectively the requests completion timeline and dataset prompt and output tokens statistics, which can be useful for debugging purpose or for deeper analysis.
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model meta-llama/Llama-3.1-8B-Instruct \
|
||||
--endpoint /v1/completions \
|
||||
--dataset-name sharegpt \
|
||||
--dataset-path <your data path>/ShareGPT_V3_unfiltered_cleaned_split.json \
|
||||
--num-prompts 100 \
|
||||
--plot-timeline \
|
||||
--timeline-itl-thresholds 2,5 \
|
||||
--plot-dataset-stats \
|
||||
--save-result
|
||||
```
|
||||
|
||||
##### Interactive Timeline
|
||||
|
||||
The generated timeline is an interactive visualization in the form of an HTML file that can be rendered in most browsers. To customize the ITL color thresholds, one can use `--timeline-itl-thresholds` flag (default: 25ms, 50ms)
|
||||
|
||||
Example output:
|
||||
|
||||
<iframe src="../../assets/contributing/vllm_bench_serve_timeline.html" width="100%" height="600" frameborder="0"></iframe>
|
||||
|
||||
##### Dataset statistics
|
||||
|
||||
The generated figure shows the input prompt and output tokens distribution.
|
||||
|
||||
Example output: 
|
||||
|
||||
#### Custom Dataset
|
||||
|
||||
If the dataset you want to benchmark is not supported yet in vLLM, even then you can benchmark on it using `CustomDataset`. Your data needs to be in `.jsonl` format and needs to have "prompt" field per entry, e.g., data.jsonl
|
||||
|
||||
+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
|
||||
|
||||
|
||||
@@ -155,9 +155,9 @@ switch to `--physcpubind=<cpu-list> --membind=<node>`.
|
||||
|
||||
These `--numa-bind*` options only apply to GPU execution processes. They do not
|
||||
configure the CPU backend's separate thread-affinity controls. Automatic
|
||||
GPU-to-NUMA detection is currently implemented for CUDA/NVML-based platforms;
|
||||
other GPU backends must provide explicit binding lists if they use these
|
||||
options.
|
||||
GPU-to-NUMA detection is currently implemented for CUDA/NVML-based as well as
|
||||
ROCM-based platforms; other GPU backends must provide explicit binding lists if
|
||||
they use these options.
|
||||
|
||||
`--numa-bind-nodes` takes one non-negative NUMA node index per visible GPU, in
|
||||
the same order as the GPU indices.
|
||||
|
||||
@@ -66,7 +66,7 @@ This is for controlling general behavior of the API when serving your model:
|
||||
|
||||
See [Audio preprocessing and chunking](#audio-preprocessing-and-chunking) for what each field controls.
|
||||
|
||||
Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server passes you the resampled waveform and task parameters; you return a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns:
|
||||
Implement the prompt construction via [get_generation_prompt][vllm.model_executor.models.interfaces.SupportsTranscription.get_generation_prompt]. The server builds a [SpeechToTextParams][vllm.config.speech_to_text.SpeechToTextParams] object that bundles the resampled waveform, task parameters, and request-specific options. Your model receives this single object and returns a valid [PromptType][vllm.inputs.llm.PromptType]. There are two common patterns:
|
||||
|
||||
#### Multimodal LLM with audio embeddings (e.g., Voxtral, Gemma3n)
|
||||
|
||||
@@ -75,21 +75,20 @@ Return a dict containing `multi_modal_data` with the audio, and either a `prompt
|
||||
??? code "get_generation_prompt()"
|
||||
|
||||
```python
|
||||
from vllm.config.speech_to_text import SpeechToTextParams
|
||||
|
||||
class YourASRModel(nn.Module, SupportsTranscription):
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def get_generation_prompt(
|
||||
cls,
|
||||
audio: np.ndarray,
|
||||
stt_config: SpeechToTextConfig,
|
||||
model_config: ModelConfig,
|
||||
language: str | None,
|
||||
task_type: Literal["transcribe", "translate"],
|
||||
request_prompt: str,
|
||||
to_language: str | None,
|
||||
stt_params: SpeechToTextParams,
|
||||
) -> PromptType:
|
||||
# Example with a free-form instruction prompt
|
||||
audio = stt_params.audio
|
||||
stt_config = stt_params.stt_config
|
||||
task_type = stt_params.task_type
|
||||
|
||||
task_word = "Transcribe" if task_type == "transcribe" else "Translate"
|
||||
prompt = (
|
||||
"<start_of_turn>user\n"
|
||||
@@ -112,20 +111,22 @@ Return a dict with separate `encoder_prompt` and `decoder_prompt` entries:
|
||||
??? code "get_generation_prompt()"
|
||||
|
||||
```python
|
||||
from vllm.config.speech_to_text import SpeechToTextParams
|
||||
|
||||
class YourASRModel(nn.Module, SupportsTranscription):
|
||||
...
|
||||
|
||||
@classmethod
|
||||
def get_generation_prompt(
|
||||
cls,
|
||||
audio: np.ndarray,
|
||||
stt_config: SpeechToTextConfig,
|
||||
model_config: ModelConfig,
|
||||
language: str | None,
|
||||
task_type: Literal["transcribe", "translate"],
|
||||
request_prompt: str,
|
||||
to_language: str | None,
|
||||
stt_params: SpeechToTextParams,
|
||||
) -> PromptType:
|
||||
audio = stt_params.audio
|
||||
stt_config = stt_params.stt_config
|
||||
language = stt_params.language
|
||||
task_type = stt_params.task_type
|
||||
request_prompt = stt_params.request_prompt
|
||||
|
||||
if language is None:
|
||||
raise ValueError("Language must be specified")
|
||||
|
||||
@@ -213,15 +214,13 @@ Relevant server logic:
|
||||
chunks = [y] if not do_split_audio else self._split_audio(y, int(sr))
|
||||
prompts = []
|
||||
for chunk in chunks:
|
||||
prompt = self.model_cls.get_generation_prompt(
|
||||
stt_params = request.build_stt_params(
|
||||
audio=chunk,
|
||||
stt_config=self.asr_config,
|
||||
model_config=self.model_config,
|
||||
language=language,
|
||||
task_type=self.task_type,
|
||||
request_prompt=request.prompt,
|
||||
to_language=to_language,
|
||||
)
|
||||
prompt = self.model_cls.get_generation_prompt(stt_params)
|
||||
prompts.append(prompt)
|
||||
return prompts, duration
|
||||
```
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ Deploying vLLM on Kubernetes is a scalable and efficient way to serve machine le
|
||||
|
||||
- [Deployment with CPUs](#deployment-with-cpus)
|
||||
- [Deployment with GPUs](#deployment-with-gpus)
|
||||
- [Serving with gRPC](#serving-with-grpc)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
- [Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"](#startup-probe-or-readiness-probe-failure-container-log-contains-keyboardinterrupt-terminated)
|
||||
- [Conclusion](#conclusion)
|
||||
@@ -387,6 +388,49 @@ INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
|
||||
|
||||
If the service is correctly deployed, you should receive a response from the vLLM model.
|
||||
|
||||
## Serving with gRPC
|
||||
|
||||
vLLM can serve models over gRPC instead of HTTP by passing the `--grpc` flag. This requires the optional gRPC dependencies:
|
||||
|
||||
```bash
|
||||
pip install vllm[grpc]
|
||||
```
|
||||
|
||||
When using `--grpc`, the server exposes the standard [gRPC Health Checking Protocol](https://github.com/grpc/grpc/blob/master/doc/health-checking.md) (`grpc.health.v1.Health`), which integrates with Kubernetes [native gRPC probes](https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/#define-a-grpc-liveness-probe) (available since Kubernetes 1.24).
|
||||
|
||||
To deploy with gRPC, change the `vllm serve` command to include `--grpc` and replace `httpGet` probes with `grpc` probes:
|
||||
|
||||
```yaml
|
||||
containers:
|
||||
- name: mistral-7b
|
||||
image: vllm/vllm-openai:latest
|
||||
command: ["/bin/sh", "-c"]
|
||||
args: [
|
||||
"pip install vllm[grpc] && vllm serve mistralai/Mistral-7B-Instruct-v0.3 --grpc --port 50051 --trust-remote-code"
|
||||
]
|
||||
ports:
|
||||
- containerPort: 50051
|
||||
livenessProbe:
|
||||
grpc:
|
||||
port: 50051
|
||||
initialDelaySeconds: 120
|
||||
periodSeconds: 10
|
||||
readinessProbe:
|
||||
grpc:
|
||||
port: 50051
|
||||
initialDelaySeconds: 120
|
||||
periodSeconds: 5
|
||||
```
|
||||
|
||||
!!! note
|
||||
The gRPC health service checks the engine status on every probe. If the engine is unhealthy or the server is shutting down, the probe returns `NOT_SERVING`.
|
||||
|
||||
You can also verify the health service manually with `grpcurl`:
|
||||
|
||||
```bash
|
||||
grpcurl -plaintext localhost:50051 grpc.health.v1.Health/Check
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Startup Probe or Readiness Probe Failure, container log contains "KeyboardInterrupt: terminated"
|
||||
|
||||
@@ -172,7 +172,7 @@ Priority is **1 = highest** (tried first).
|
||||
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
@@ -213,7 +213,7 @@ configuration.
|
||||
| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
|
||||
| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |
|
||||
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 512, 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |
|
||||
| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |
|
||||
| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
|
||||
@@ -44,7 +44,7 @@ The table below lists the quantization schemes supported by each fusion on each
|
||||
| `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — |
|
||||
| `fuse_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — |
|
||||
| `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* |
|
||||
| `fuse_attn_quant` (MLA)\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static(untested)\* |
|
||||
| `fuse_attn_quant` (MLA)\* | FP8 static\*, FP8 per-group\*, NVFP4\* | FP8 static\*, FP8 per-group\* | FP8 static\*, FP8 per-group\* | — | FP8 static\* (untested) |
|
||||
| `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |
|
||||
| `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16† | FP16/BF16† | — |
|
||||
| `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — |
|
||||
@@ -152,7 +152,7 @@ standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patt
|
||||
|
||||
- `FLASHINFER`: CUDA sm100+ with FlashInfer installed
|
||||
|
||||
`MLAAttention → FP8 static quant` / `MLAAttention → NVFP4 dynamic quant`:
|
||||
`MLAAttention → FP8 static, FP8 per-group, NVFP4 dynamic quant`
|
||||
|
||||
The MLA fusion operates at the graph level on the `unified_mla_attention_with_output` op and works
|
||||
with all MLA decode and prefill backend combinations. Unlike standard `Attention` backends (where
|
||||
|
||||
@@ -36,7 +36,7 @@ th {
|
||||
| deepep_high_throughput | standard | fp8 | G(128),A,T<sup>2</sup> | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] |
|
||||
| deepep_low_latency | batched | fp8 | G(128),A,T<sup>3</sup> | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] |
|
||||
|
||||
!!! info "Table key"
|
||||
1. All types: mxfp4, nvfp4, int4, int8, fp8
|
||||
@@ -83,8 +83,8 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
|
||||
| triton | standard | all<sup>1</sup> | G,A,T | silu, gelu,</br>swigluoai,</br>silu_no_mul,</br>gelu_no_mul | Y | Y | [`fused_experts`][vllm.model_executor.layers.fused_moe.fused_moe.fused_experts],</br>[`TritonExperts`][vllm.model_executor.layers.fused_moe.fused_moe.TritonExperts] |
|
||||
| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
|
||||
| deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | </br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
|
||||
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
|
||||
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
|
||||
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp4] |
|
||||
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.experts.cutlass_moe.CutlassBatchedExpertsFp8] |
|
||||
| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
|
||||
| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
|
||||
| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
# Context Extension
|
||||
|
||||
!!! note
|
||||
The `--rope-scaling` parameter used in older versions of vLLM is no longer supported. Please use the `--hf-overrides` method with `rope_parameters` instead.
|
||||
This directory contains examples for extending the context length of models using vLLM.
|
||||
|
||||
## Offline Inference 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/features/context_extension/context_extension_offline.py
|
||||
```
|
||||
|
||||
## OpenAI Online Method
|
||||
|
||||
You can also use vLLM's OpenAI-compatible API to serve models with extended context length.
|
||||
|
||||
### Usage
|
||||
|
||||
Run the vLLM server with the following command to extend the context length using YARN:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-0.6B \
|
||||
--hf-overrides '{"rope_parameters": {"factor": 4.0, "original_max_position_embeddings": 32768, "rope_theta": 1000000, "rope_type": "yarn"}}' \
|
||||
--max-model-len 131072
|
||||
```
|
||||
|
||||
### Client Example
|
||||
|
||||
After starting the server, you can use the OpenAI Python client to interact with it:
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(
|
||||
base_url="http://localhost:8000/v1",
|
||||
api_key="token-abc123" # Dummy API key, required by the client
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant"},
|
||||
{"role": "user", "content": "Hello"}
|
||||
],
|
||||
max_tokens=128,
|
||||
temperature=0.8,
|
||||
top_p=0.95
|
||||
)
|
||||
|
||||
print(response.choices[0].message.content)
|
||||
```
|
||||
|
||||
### Key Parameters
|
||||
|
||||
The available parameters depend on the `rope_type` you choose. For detailed information about all supported RoPE types and their specific parameters, please refer to the [Hugging Face Transformers RoPE documentation](https://huggingface.co/docs/transformers/main/en/internal/rope_utils#transformers.RopeParameters).
|
||||
|
||||
Common parameters include:
|
||||
|
||||
- `rope_type`: The type of RoPE implementation (e.g., "yarn", "linear", "dynamic")
|
||||
- `factor`: The factor by which to extend the context length
|
||||
- `original_max_position_embeddings`: The original maximum position embeddings of the model
|
||||
|
||||
The following parameters are specific to vLLM:
|
||||
|
||||
- `max_model_len`: The new maximum sequence length after extension (original * factor).
|
||||
Used for KV cache pre‑allocation and request limit at serving time.
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -68,7 +68,7 @@ You can pass a single image to the `'image'` field of the multi-modal dictionary
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
|
||||
|
||||
To substitute multiple images inside the same text prompt, you can pass in a list of images instead:
|
||||
|
||||
@@ -101,7 +101,7 @@ To substitute multiple images inside the same text prompt, you can pass in a lis
|
||||
print(generated_text)
|
||||
```
|
||||
|
||||
Full example: [examples/offline_inference/vision_language_multi_image.py](../../examples/offline_inference/vision_language_multi_image.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_multi_image_offline.py](../../examples/generate/multimodal/vision_language_multi_image_offline.py)
|
||||
|
||||
If using the [LLM.chat](../models/generative_models.md#llmchat) method, you can pass images directly in the message content using various formats: image URLs, PIL Image objects, or pre-computed embeddings:
|
||||
|
||||
@@ -287,13 +287,13 @@ Instead of NumPy arrays, you can also pass `'torch.Tensor'` instances, as shown
|
||||
!!! note
|
||||
'process_vision_info' is only applicable to Qwen2.5-VL and similar models.
|
||||
|
||||
Full example: [examples/offline_inference/vision_language.py](../../examples/offline_inference/vision_language.py)
|
||||
Full example: [examples/generate/multimodal/vision_language_offline.py](../../examples/generate/multimodal/vision_language_offline.py)
|
||||
|
||||
### Audio Inputs
|
||||
|
||||
You can pass a tuple `(array, sampling_rate)` to the `'audio'` field of the multi-modal dictionary.
|
||||
|
||||
Full example: [examples/offline_inference/audio_language.py](../../examples/offline_inference/audio_language.py)
|
||||
Full example: [examples/generate/multimodal/audio_language_offline.py](../../examples/generate/multimodal/audio_language_offline.py)
|
||||
|
||||
#### Chunking Long Audio for Transcription
|
||||
|
||||
@@ -674,7 +674,7 @@ Then, you can use the OpenAI client as follows:
|
||||
print("Chat completion output:", chat_response.choices[0].message.content)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! tip
|
||||
Loading from local file paths is also supported on vLLM: You can specify the allowed local media path via `--allowed-local-media-path` when launching the API server/engine,
|
||||
@@ -745,7 +745,7 @@ Then, you can use the OpenAI client as follows:
|
||||
print("Chat completion output from image url:", result)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! note
|
||||
By default, the timeout for fetching videos through HTTP URL is `30` seconds.
|
||||
@@ -780,6 +780,70 @@ vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \
|
||||
|
||||
Works with common video formats like MP4 when using OpenCV backends.
|
||||
|
||||
#### Pre-extracted Frame Sequences with `media_io_kwargs`
|
||||
|
||||
When you extract video frames on the client side and send them as `video/jpeg` (base64-concatenated JPEG frames), you can preserve the original video metadata by using `media_io_kwargs` in your request. This enables more accurate video understanding by preserving temporal information that would otherwise be lost during client-side frame extraction.
|
||||
|
||||
**Supported Parameters:**
|
||||
|
||||
| Parameter | Type | Description |
|
||||
| --------- | ---- | ----------- |
|
||||
| `fps` | float | Frame rate of the original video |
|
||||
| `frames_indices` | list[int] | Indices of the actually sampled frames |
|
||||
| `total_num_frames` | int | Total frame count of the original video |
|
||||
| `duration` | float | Duration of the original video in seconds |
|
||||
| `do_sample_frames` | bool | Whether to perform frame sampling |
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8000/v1", api_key="EMPTY")
|
||||
|
||||
# Client-side frame extraction
|
||||
frames = extract_frames(video_path, num_frames=32)
|
||||
frames_b64 = ",".join([encode_image(f) for f in frames])
|
||||
video_url = f"data:video/jpeg;base64,{frames_b64}"
|
||||
|
||||
# Pass video metadata via media_io_kwargs
|
||||
response = client.chat.completions.create(
|
||||
model="your-multimodal-model",
|
||||
messages=[{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "video_url", "video_url": {"url": video_url}},
|
||||
{"type": "text", "text": "Describe what happens in this video."}
|
||||
]
|
||||
}],
|
||||
extra_body={
|
||||
"media_io_kwargs": {
|
||||
"video": {
|
||||
"fps": 30.0,
|
||||
"frames_indices": [0, 10, 20, 30, 40, 50, 60, 70, 80, 90,
|
||||
100, 110, 120, 130, 140, 150, 160, 170,
|
||||
180, 190, 200, 210, 220, 230, 240, 250,
|
||||
260, 270, 280, 290, 300, 310],
|
||||
"total_num_frames": 900,
|
||||
"duration": 30.0,
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
print(response.choices[0].message.content)
|
||||
```
|
||||
|
||||
**Why use `media_io_kwargs`?**
|
||||
|
||||
When extracting frames client-side, the server loses important context about the original video:
|
||||
|
||||
- **Temporal information**: Which frames were sampled and their positions in the original timeline
|
||||
- **Video duration**: How long the original video was
|
||||
- **Frame rate**: The original playback speed
|
||||
|
||||
By passing this metadata, the model can better understand the temporal distribution of the sampled frames and whether important moments might have been skipped.
|
||||
|
||||
#### Custom RGBA Background Color
|
||||
|
||||
To use a custom background color for RGBA images, pass the `rgba_background_color` parameter via `--media-io-kwargs`:
|
||||
@@ -894,7 +958,7 @@ Alternatively, you can pass `audio_url`, which is the audio counterpart of `imag
|
||||
print("Chat completion output from audio url:", result)
|
||||
```
|
||||
|
||||
Full example: [examples/online_serving/openai_chat_completion_client_for_multimodal.py](../../examples/online_serving/openai_chat_completion_client_for_multimodal.py)
|
||||
Full example: [examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py](../../examples/generate/multimodal/openai_chat_completion_client_for_multimodal.py)
|
||||
|
||||
!!! note
|
||||
By default, the timeout for fetching audios through HTTP URL is `10` seconds.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -20,6 +20,7 @@ The following are the supported quantization formats for vLLM:
|
||||
- [AMD Quark](quark.md)
|
||||
- [Quantized KV Cache](quantized_kvcache.md)
|
||||
- [TorchAO](torchao.md)
|
||||
- [FP8 ViT Encoder Attention](fp8_vit_attn.md)
|
||||
|
||||
## Supported Hardware
|
||||
|
||||
|
||||
@@ -0,0 +1,109 @@
|
||||
# FP8 ViT Encoder Attention
|
||||
|
||||
For visual understanding workloads with large images (e.g. QHD, 4K) and relatively
|
||||
short text prompts/generation, the ViT encoder attention can become a significant
|
||||
bottleneck, especially when the text model is quantized (e.g. NVFP4). vLLM
|
||||
supports optional FP8 quantization for the ViT encoder attention via the
|
||||
FlashInfer cuDNN backend. Q/K/V are quantized on-the-fly to FP8 before the
|
||||
cuDNN attention call.
|
||||
|
||||
!!! note
|
||||
- Currently supports Qwen3-VL family models only (`qwen3_vl`, `qwen3_vl_moe`,
|
||||
`qwen3_5`, `qwen3_5_moe`, and other models using Qwen3 ViT).
|
||||
- Dynamic scaling is not compatible with ViT full CUDA graphs.
|
||||
- Performance gains are mostly visible at QHD/4K resolutions or multi-image
|
||||
requests. Smaller images may see no speedup due to quantization overhead
|
||||
(3 quantization kernel launches + un-padding).
|
||||
- FP8 tensor-core speedup is more pronounced on GB300 than GB200.
|
||||
|
||||
## Requirements
|
||||
|
||||
- FlashInfer cuDNN backend with cuDNN >= 9.17.1.
|
||||
|
||||
## Usage
|
||||
|
||||
Enable FP8 ViT attention by passing `--mm-encoder-attn-dtype fp8` together
|
||||
with `--mm-encoder-attn-backend FLASHINFER`:
|
||||
|
||||
```bash
|
||||
vllm serve $MODEL \
|
||||
--mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8
|
||||
```
|
||||
|
||||
By default (no scale file), **dynamic scaling** is used: a 16-entry circular
|
||||
buffer of observed Q/K/V amax values drives per-forward scale updates. This
|
||||
matches BF16 accuracy without any calibration but adds a small per-forward
|
||||
overhead.
|
||||
|
||||
## Calibrate-Once, Reuse Workflow (Recommended)
|
||||
|
||||
For production, calibrate static scales on a representative dataset once and
|
||||
reuse them to avoid the dynamic overhead:
|
||||
|
||||
```bash
|
||||
# Step 1: calibrate and save scales (runs dynamic scaling for 16 passes,
|
||||
# then dumps the learned scales to JSON).
|
||||
vllm bench mm-processor \
|
||||
--model $MODEL --mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8 \
|
||||
--mm-encoder-fp8-scale-save-path /path/to/scales.json \
|
||||
--dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \
|
||||
--num-prompts 100
|
||||
|
||||
# Step 2: serve with static scales (no dynamic overhead).
|
||||
vllm serve $MODEL \
|
||||
--mm-encoder-attn-backend FLASHINFER \
|
||||
--mm-encoder-attn-dtype fp8 \
|
||||
--mm-encoder-fp8-scale-path /path/to/scales.json
|
||||
```
|
||||
|
||||
Saved scales are multiplied by `--mm-encoder-fp8-scale-save-margin` (default
|
||||
`1.5`) to leave headroom against activation outliers not present in the
|
||||
calibration set. The default has been validated to generalize across datasets
|
||||
(e.g. VisionArena-Chat calibration maintains BF16 accuracy on ChartQA).
|
||||
|
||||
## Scale File Format
|
||||
|
||||
```json
|
||||
{
|
||||
"visual.blocks.0.attn.attn": {"q": 224.0, "k": 198.0, "v": 210.0},
|
||||
"visual.blocks.1.attn.attn": {"q": 218.0, "k": 195.0, "v": 207.0}
|
||||
}
|
||||
```
|
||||
|
||||
Keys `q_scale` / `k_scale` / `v_scale` are accepted as aliases.
|
||||
|
||||
## Performance
|
||||
|
||||
**Core cuDNN attention kernel** (PyTorch profiler, `cudnn_generated_fort_native_sdpa_sm100_flash_fprop`, head_dim=128, seq_len=8192):
|
||||
|
||||
| Hardware | BF16 | FP8 | Speedup |
|
||||
| -------- | ---- | ---- | ------- |
|
||||
| GB200 | 350 us | 312 us | **1.12x** |
|
||||
| GB300 | 300 us | 211 us | **1.42x** |
|
||||
|
||||
**End-to-end encoder forward time** (Qwen3-VL-30B-A3B-Instruct on GB200, 3 images/request):
|
||||
|
||||
| Resolution | BF16 median | FP8 median | Speedup |
|
||||
| ---------- | ----------- | ---------- | ------- |
|
||||
| HD (720x1280) | 31.77 ms | 36.39 ms | 0.87x |
|
||||
| FullHD (1080x1920) | 57.99 ms | 58.73 ms | ~same |
|
||||
| QHD (1440x2560) | 131.83 ms | 122.30 ms | **1.08x** |
|
||||
| 4K (2160x3840) | 543.44 ms | 460.31 ms | **1.18x** |
|
||||
|
||||
Crossover is around FullHD with 3 images/request. At QHD and above, FP8 wins.
|
||||
|
||||
## Accuracy
|
||||
|
||||
ChartQA, Qwen3-VL-8B-Instruct, 500 samples. FP8 static uses scales calibrated
|
||||
on VisionArena-Chat (with default 1.5x margin):
|
||||
|
||||
| Metric | BF16 | FP8 dynamic | FP8 static |
|
||||
| ------ | ---- | ----------- | ---------- |
|
||||
| relaxed_accuracy | 0.780 | 0.776 | 0.780 |
|
||||
| anywhere_accuracy | 0.806 | 0.816 | 0.814 |
|
||||
| exact_match | 0.584 | 0.582 | 0.578 |
|
||||
|
||||
All three configurations match within statistical noise, confirming that
|
||||
static scales calibrated on one dataset generalize to another.
|
||||
@@ -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` | ❌ |
|
||||
@@ -202,7 +203,7 @@ The reasoning content is also available when both tool calling and the reasoning
|
||||
print(f"Arguments: {tool_call.arguments}")
|
||||
```
|
||||
|
||||
For more examples, please refer to [examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/online_serving/openai_chat_completion_tool_calls_with_reasoning.py).
|
||||
For more examples, please refer to [examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py](../../examples/reasoning/openai_chat_completion_tool_calls_with_reasoning.py).
|
||||
|
||||
## Server-Level Default Chat Template Kwargs
|
||||
|
||||
|
||||
@@ -32,9 +32,102 @@ 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
|
||||
|
||||
Use `--speculative-config` to pass speculative decoding settings as a JSON
|
||||
object on the CLI:
|
||||
|
||||
```bash
|
||||
vllm serve <target-model> \
|
||||
--speculative-config '{
|
||||
"method": "draft_model",
|
||||
"model": "<draft-model>",
|
||||
"num_speculative_tokens": 5
|
||||
}'
|
||||
```
|
||||
|
||||
The same keys are accepted from Python via `LLM(..., speculative_config={...})`.
|
||||
The tables below highlight common user-facing keys accepted in this JSON
|
||||
object; they are not an exhaustive schema reference.
|
||||
For more details, see the generated [engine arguments reference](../../configuration/engine_args.md)
|
||||
and the API docs for [vllm.config.SpeculativeConfig][].
|
||||
|
||||
### Common keys
|
||||
|
||||
These keys are commonly used across speculative decoding setups, though some
|
||||
only apply to model-based methods such as `draft_model`, `mtp`, `eagle3`, and
|
||||
`dflash`.
|
||||
|
||||
| Key | Type | Default | Allowed values / meaning |
|
||||
| --- | --- | --- | --- |
|
||||
| `method` | `string` | `None` | Speculation method. Common values include `draft_model`, `ngram`, `suffix`, `mtp`, `eagle3`, and `dflash`. If omitted, vLLM infers the method from the provided configuration when possible. |
|
||||
| `model` | `string` | `None` | Draft model, EAGLE head, or auxiliary model identifier. For `ngram`, `ngram_gpu`, `suffix`, and `mtp`, this can often be omitted. |
|
||||
| `num_speculative_tokens` | `integer > 0` | `None` | Number of speculative tokens to propose per step. Required for methods that do not infer it from model metadata. |
|
||||
| `draft_tensor_parallel_size` | `integer >= 1` | `None` | Tensor parallel size for the draft model. |
|
||||
| `max_model_len` | `integer >= 1` | `None` | Maximum context length for the draft model. |
|
||||
| `parallel_drafting` | `boolean` | `false` | Enable parallel draft token generation. Only compatible with EAGLE and draft-model methods. |
|
||||
| `rejection_sample_method` | `string` | `strict` | `strict`, `probabilistic`, or `synthetic`. |
|
||||
| `synthetic_acceptance_rate` | `float` | `None` | Average acceptance rate to target when `rejection_sample_method` is `synthetic`. Valid range is `[0, 1]`. |
|
||||
|
||||
### Method-specific keys
|
||||
|
||||
#### N-gram
|
||||
|
||||
| Key | Type | Default | Meaning |
|
||||
| --- | --- | --- | --- |
|
||||
| `prompt_lookup_max` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_min` when omitted | Maximum n-gram window size. |
|
||||
| `prompt_lookup_min` | `integer >= 1` | `5` if both lookup bounds are omitted; otherwise mirrors `prompt_lookup_max` when omitted | Minimum n-gram window size. |
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
vllm serve <target-model> \
|
||||
--speculative-config '{
|
||||
"method": "ngram",
|
||||
"num_speculative_tokens": 4,
|
||||
"prompt_lookup_min": 2,
|
||||
"prompt_lookup_max": 5
|
||||
}'
|
||||
```
|
||||
|
||||
#### Suffix decoding
|
||||
|
||||
| Key | Type | Default | Meaning |
|
||||
| --- | --- | --- | --- |
|
||||
| `suffix_decoding_max_tree_depth` | `integer` | `24` | Maximum combined prefix-match and speculation tree depth. |
|
||||
| `suffix_decoding_max_cached_requests` | `integer` | `10000` | Maximum number of requests cached in the global suffix tree. Set `0` to disable the global cache. |
|
||||
| `suffix_decoding_max_spec_factor` | `float` | `1.0` | Caps speculative length as a multiple of prefix-match length. |
|
||||
| `suffix_decoding_min_token_prob` | `float` | `0.1` | Minimum estimated token probability required to speculate a token. |
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
vllm serve <target-model> \
|
||||
--speculative-config '{
|
||||
"method": "suffix",
|
||||
"num_speculative_tokens": 8,
|
||||
"suffix_decoding_max_tree_depth": 24,
|
||||
"suffix_decoding_max_cached_requests": 10000,
|
||||
"suffix_decoding_max_spec_factor": 1.0,
|
||||
"suffix_decoding_min_token_prob": 0.1
|
||||
}'
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- `--speculative-config` expects a JSON object on the CLI. In YAML config
|
||||
files, use a nested mapping instead of an escaped JSON string.
|
||||
- `tensor_parallel_size` is not a valid key in `speculative_config`. Use
|
||||
`draft_tensor_parallel_size` instead.
|
||||
- Keys such as `temperature` and `top_p` are sampling parameters, not
|
||||
`--speculative-config` fields.
|
||||
- Internal fields such as `target_model_config`, `draft_model_config`,
|
||||
`target_parallel_config`, `draft_parallel_config`, and `draft_load_config`
|
||||
are populated by vLLM and are not intended to be set by users.
|
||||
|
||||
## Lossless guarantees of Speculative Decoding
|
||||
|
||||
In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
|
||||
|
||||
@@ -33,9 +33,9 @@ vllm serve Qwen/Qwen3-4B-Thinking-2507 \
|
||||
--port 8000 \
|
||||
--seed 42 \
|
||||
-tp 1 \
|
||||
--max_model_len 2048 \
|
||||
--gpu_memory_utilization 0.8 \
|
||||
--speculative_config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
|
||||
--max-model-len 2048 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--speculative-config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
|
||||
```
|
||||
|
||||
The code used to request as completions as a client remains unchanged:
|
||||
@@ -77,4 +77,8 @@ The code used to request as completions as a client remains unchanged:
|
||||
```
|
||||
|
||||
!!! warning
|
||||
Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated.
|
||||
Note: Please use `--speculative-config` to set all configurations related
|
||||
to speculative decoding. The previous method of specifying the model
|
||||
through `--speculative-model` and adding related parameters such as
|
||||
`--num-speculative-tokens` separately has been deprecated. For supported
|
||||
keys and examples, see the [`--speculative-config` schema](README.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
|
||||
|
||||
|
||||
@@ -38,7 +38,7 @@ for output in outputs:
|
||||
```bash
|
||||
vllm serve XiaomiMiMo/MiMo-7B-Base \
|
||||
--tensor-parallel-size 1 \
|
||||
--speculative_config '{"method":"mtp","num_speculative_tokens":1}'
|
||||
--speculative-config '{"method":"mtp","num_speculative_tokens":1}'
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
@@ -36,9 +36,9 @@ vllm serve Qwen/Qwen3-4B \
|
||||
--port 8000 \
|
||||
--seed 42 \
|
||||
-tp 1 \
|
||||
--max_model_len 2048 \
|
||||
--gpu_memory_utilization 0.8 \
|
||||
--speculative_config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}'
|
||||
--max-model-len 2048 \
|
||||
--gpu-memory-utilization 0.8 \
|
||||
--speculative-config '{"model": "amd/PARD-Qwen3-0.6B", "num_speculative_tokens": 12, "method": "draft_model", "parallel_drafting": true}'
|
||||
```
|
||||
|
||||
## Pre-trained PARD weights
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -375,8 +375,8 @@ For (G)B300, we recommend using CUDA 13, as shown in the following command.
|
||||
|
||||
```bash
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--build-arg CUDA_VERSION=13.0.1 \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.1-devel-ubuntu22.04 \
|
||||
--build-arg CUDA_VERSION=13.0.2 \
|
||||
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
|
||||
--build-arg max_jobs=256 \
|
||||
--build-arg nvcc_threads=2 \
|
||||
--build-arg RUN_WHEEL_CHECK=false \
|
||||
|
||||
@@ -46,7 +46,7 @@ pip install -v -r requirements/xpu.txt
|
||||
|
||||
!!! note
|
||||
- `triton` (without suffix) is for NVIDIA GPUs only. On XPU, using it instead of `triton-xpu` can cause correctness or runtime issues.
|
||||
- For torch 2.10 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.6.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
|
||||
- For torch 2.11 (the version used in `requirements/xpu.txt`), the matching package is `triton-xpu==3.7.0`. If you use a different version of torch, check the corresponding `triton-xpu` version in [docker/Dockerfile.xpu](https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile.xpu).
|
||||
|
||||
- Finally, build and install vLLM XPU backend:
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
|
||||
@@ -78,7 +78,7 @@ The scoring models is designed to compute similarity scores between two input pr
|
||||
|-----------------------|---------------|----------------------------------------------|--------------------|--------------------------|
|
||||
| `classify` (see note) | Sequence-wise | reranker score for each sequence | `cross-encoder` | linear classifier |
|
||||
| `embed` | Sequence-wise | vector representations for each sequence | `bi-encoder` | cosine similarity |
|
||||
| `token_classify` | Token-wise | probability vector of classes for each token | nan | nan |
|
||||
| `token_classify` | Token-wise | probability vector of classes for each token | N/A | N/A |
|
||||
| `token_embed` | Token-wise | vector representations for each token | `late-interaction` | late interaction(MaxSim) |
|
||||
|
||||
!!! note
|
||||
@@ -86,14 +86,15 @@ The scoring models is designed to compute similarity scores between two input pr
|
||||
|
||||
### Pooling Usages
|
||||
|
||||
| Pooling Usages | Description |
|
||||
|-----------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| Classification Usages | Predicting which predefined category, class, or label best corresponds to a given input. |
|
||||
| Embedding Usages | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings). |
|
||||
| Token Classification Usages | Token-wise classification |
|
||||
| Token Embedding Usages | Token-wise embedding |
|
||||
| Scoring Usages | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. |
|
||||
| Reward Usages | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences. |
|
||||
| Pooling Usages | Description |
|
||||
|-----------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| Classification Usages | Predicting which predefined category, class, or label best corresponds to a given input. |
|
||||
| Embedding Usages | Converts unstructured data (text, images, audio, etc.) into structured numerical vectors (embeddings). |
|
||||
| Token Classification Usages | Token-wise classification |
|
||||
| Token Embedding Usages | Token-wise embedding |
|
||||
| Reward Usages | Evaluates the quality of outputs generated by a language model, acting as a proxy for human preferences. |
|
||||
| Scoring Usages | Computes similarity scores between two inputs. It supports three model types (aka `score_type`): `cross-encoder`, `late-interaction`, and `bi-encoder`. |
|
||||
| Plugins Usages | Allow users to customize input and output processors. For more information, please refer to [IO Processor Plugins](../../design/io_processor_plugins.md). |
|
||||
|
||||
We also have some special models that support multiple pooling tasks, or have specific usage scenarios, or support special inputs and outputs.
|
||||
|
||||
@@ -101,9 +102,9 @@ For more detailed information, please refer to the link below.
|
||||
|
||||
- [Classification Usages](classify.md)
|
||||
- [Embedding Usages](embed.md)
|
||||
- [Reward Usages](reward.md)
|
||||
- [Token Classification Usages](token_classify.md)
|
||||
- [Token Embedding Usages](token_embed.md)
|
||||
- [Reward Usages](reward.md)
|
||||
- [Scoring Usages](scoring.md)
|
||||
- [Specific Model Examples](specific_models.md)
|
||||
|
||||
@@ -113,15 +114,17 @@ Each pooling model in vLLM supports one or more of these tasks according to
|
||||
[Pooler.get_supported_tasks][vllm.model_executor.layers.pooler.Pooler.get_supported_tasks],
|
||||
enabling the corresponding APIs.
|
||||
|
||||
### Offline APIs corresponding to pooling tasks
|
||||
### Offline APIs corresponding to pooling usages
|
||||
|
||||
| Task | APIs |
|
||||
|------------------|---------------------------------------------------------------------------------------|
|
||||
| `embed` | `LLM.embed(...)`, `LLM.encode(..., pooling_task="embed")`, `LLM.score(...)`(see note) |
|
||||
| `classify` | `LLM.classify(...)`, `LLM.encode(..., pooling_task="classify")`, `LLM.score(...)` |
|
||||
| `token_classify` | `LLM.reward(...)`, `LLM.encode(..., pooling_task="token_classify")` |
|
||||
| `token_embed` | `LLM.encode(..., pooling_task="token_embed")`, `LLM.score(...)` |
|
||||
| `plugin` | `LLM.encode(..., pooling_task="plugin")` |
|
||||
| Pooling Usages | Dedicated API | Pooling task for `LLM.encode` API | Score Types | scoring function |
|
||||
|-----------------------------|---------------------|-----------------------------------|----------------------------|--------------------------|
|
||||
| Classification Usages | `LLM.classify(...)` | `classify` | `cross-encoder` (see note) | linear classifier |
|
||||
| Embedding Usages | `LLM.embed(...)` | `embed` | `bi-encoder` | cosine similarity |
|
||||
| Token Classification Usages | N/A | `token_classify` | N/A | N/A |
|
||||
| Token Embedding Usages | N/A | `token_embed` | `late-interaction` | late interaction(MaxSim) |
|
||||
| Reward Usages | N/A | `classify` & `token_classify` | N/A | N/A |
|
||||
| Scoring Usages | `LLM.score(...)` | N/A | N/A | N/A |
|
||||
| Plugins Usages | N/A | `plugin` | N/A | N/A |
|
||||
|
||||
!!! note
|
||||
Only when a classification model outputs num_labels equal to 1 can it be used as a scoring model and have its scoring API enabled.
|
||||
@@ -147,7 +150,7 @@ It is primarily designed for [score models](scoring.md).
|
||||
|
||||
The [encode][vllm.LLM.encode] method is available to all pooling models in vLLM.
|
||||
|
||||
Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
|
||||
Please use one of the more specific methods or set the task directly when using `LLM.encode`, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).
|
||||
|
||||
### Examples
|
||||
|
||||
@@ -183,9 +186,12 @@ Our Pooling API (`/pooling`) is similar to `LLM.encode`, being applicable to all
|
||||
|
||||
The input format is the same as [Embeddings API](embed.md#openai-compatible-embeddings-api), but the output data can contain an arbitrary nested list, not just a 1-D list of floats.
|
||||
|
||||
Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-tasks).
|
||||
Please use one of the more specific APIs or set the task directly when using the Pooling API, refer to the [table above](#offline-apis-corresponding-to-pooling-usages).
|
||||
|
||||
Code example: [examples/pooling/pooling/pooling_online.py](../../../examples/pooling/pooling/pooling_online.py)
|
||||
Code examples:
|
||||
|
||||
- [Online example](../../../examples/pooling/reward/token_reward_online.py)
|
||||
- [Offline example](../../../examples/pooling/reward/token_reward_offline.py)
|
||||
|
||||
### Examples
|
||||
|
||||
|
||||
@@ -134,3 +134,13 @@ print(f"Data: {data!r}")
|
||||
## Online Serving
|
||||
|
||||
Please refer to the [pooling API](README.md#pooling-api). Pooling task corresponding to reward model types refer to the [table above](#summary).
|
||||
|
||||
## More examples
|
||||
|
||||
More examples can be found here: [examples/pooling/reward](../../../examples/pooling/reward)
|
||||
|
||||
## Deprecated Features
|
||||
|
||||
### `LLM.reward`
|
||||
|
||||
`llm.reward` api is deprecated and will be removed in v0.23. Please use `LLM.encode` with `pooling_task="classify"` or `pooling_task="token_classify"` instead.
|
||||
|
||||
@@ -384,6 +384,7 @@ th {
|
||||
| `DeepseekForCausalLM` | DeepSeek | `deepseek-ai/deepseek-llm-67b-base`, `deepseek-ai/deepseek-llm-7b-chat`, etc. | ✅︎ | ✅︎ |
|
||||
| `DeepseekV2ForCausalLM` | DeepSeek-V2 | `deepseek-ai/DeepSeek-V2`, `deepseek-ai/DeepSeek-V2-Chat`, etc. | ✅︎ | ✅︎ |
|
||||
| `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`, etc. | ✅︎ | ✅︎ |
|
||||
| `DeepseekV4ForCausalLM` | DeepSeek-V4 | `deepseek-ai/DeepSeek-V4-Flash`, `deepseek-ai/DeepSeek-V4-Pro`, etc. | | |
|
||||
| `Dots1ForCausalLM` | dots.llm1 | `rednote-hilab/dots.llm1.base`, `rednote-hilab/dots.llm1.inst`, etc. | | ✅︎ |
|
||||
| `DotsOCRForCausalLM` | dots_ocr | `rednote-hilab/dots.ocr` | ✅︎ | ✅︎ |
|
||||
| `Ernie4_5ForCausalLM` | Ernie4.5 | `baidu/ERNIE-4.5-0.3B-PT`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -419,6 +420,7 @@ th {
|
||||
| `Grok1ForCausalLM` | Grok2 | `xai-org/grok-2` | ✅︎ | ✅︎ |
|
||||
| `HunYuanDenseV1ForCausalLM` | Hunyuan Dense | `tencent/Hunyuan-7B-Instruct` | ✅︎ | ✅︎ |
|
||||
| `HunYuanMoEV1ForCausalLM` | Hunyuan-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | ✅︎ |
|
||||
| `HYV3ForCausalLM` | HY3 | `tencent/Hy3-preview-Base`, `tencent/Hy3-preview` | ✅︎ | ✅︎ |
|
||||
| `HyperCLOVAXForCausalLM` | HyperCLOVAX-SEED-Think-14B | `naver-hyperclovax/HyperCLOVAX-SEED-Think-14B` | ✅︎ | ✅︎ |
|
||||
| `InternLMForCausalLM` | InternLM | `internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
| `InternLM2ForCausalLM` | InternLM2 | `internlm/internlm2-7b`, `internlm/internlm2-chat-7b`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -437,6 +439,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. | | ✅︎ |
|
||||
| `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. | | |
|
||||
@@ -472,6 +475,7 @@ th {
|
||||
| `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B`, etc. | ✅︎ | ✅︎ |
|
||||
| `Qwen3NextForCausalLM` | Qwen3NextMoE | `Qwen/Qwen3-Next-80B-A3B-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
| `RWForCausalLM` | Falcon RW | `tiiuae/falcon-40b`, etc. | | ✅︎ |
|
||||
| `Rnj1ForCausalLM` | Rnj1 | `EssentialAI/rnj-1-instruct`, etc. | | |
|
||||
| `SarvamMoEForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-30b-a3b`, etc. | ✅︎ | ✅︎ |
|
||||
| `SarvamMLAForCausalLM` | Sarvam 2 | `sarvamai/sarvam2-105b-a9b`, etc. | | ✅︎ |
|
||||
| `SeedOssForCausalLM` | SeedOss | `ByteDance-Seed/Seed-OSS-36B-Instruct`, etc. | ✅︎ | ✅︎ |
|
||||
@@ -587,6 +591,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen
|
||||
| `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |
|
||||
| `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I<sup>+</sup> + V<sup>+</sup> | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |
|
||||
| `MiDashengLMModel` | MiDashengLM | T + A<sup>+</sup> | `mispeech/midashenglm-7b` | | ✅︎ |
|
||||
| `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>+</sup> | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ |
|
||||
| `MiniCPMO` | MiniCPM-O | T + I<sup>E+</sup> + V<sup>E+</sup> + A<sup>E+</sup> | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |
|
||||
| `MiniCPMV` | MiniCPM-V | T + I<sup>E+</sup> + V<sup>E+</sup> | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | |
|
||||
| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + I<sup>E+</sup> | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |
|
||||
@@ -641,10 +646,10 @@ Some models are supported only via the [Transformers modeling backend](#transfor
|
||||
!!! note
|
||||
`Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its
|
||||
MobileNet-v5 vision backbone.
|
||||
|
||||
|
||||
Performance is not yet fully optimized mainly due to:
|
||||
|
||||
- Both audio and vision MM encoders use `transformers.AutoModel` implementation.
|
||||
|
||||
- Both audio and vision MM encoders use `transformers.AutoModel` implementation.
|
||||
- There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.
|
||||
|
||||
!!! note
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -251,7 +251,7 @@ The following extra parameters are supported:
|
||||
Our Responses API is compatible with [OpenAI's Responses API](https://platform.openai.com/docs/api-reference/responses);
|
||||
you can use the [official OpenAI Python client](https://github.com/openai/openai-python) to interact with it.
|
||||
|
||||
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/online_serving/openai_responses_client_with_tools.py)
|
||||
Code example: [examples/online_serving/openai_responses_client_with_tools.py](../../examples/tool_calling/openai_responses_client_with_tools.py)
|
||||
|
||||
#### Extra parameters
|
||||
|
||||
@@ -279,7 +279,7 @@ you can use the [official OpenAI Python client](https://github.com/openai/openai
|
||||
!!! note
|
||||
To use the Transcriptions API, please install with extra audio dependencies using `pip install vllm[audio]`.
|
||||
|
||||
Code example: [examples/online_serving/openai_transcription_client.py](../../examples/online_serving/openai_transcription_client.py)
|
||||
Code example: [examples/speech_to_text/openai/openai_transcription_client.py](../../examples/speech_to_text/openai/openai_transcription_client.py)
|
||||
|
||||
NOTE: beam search is currently supported in the transcriptions endpoint for encoder-decoder multimodal models, e.g., whisper, but highly inefficient as work for handling the encoder/decoder cache is actively ongoing. This is an active point of ongoing optimization and will be handled properly in the very near future.
|
||||
|
||||
@@ -397,7 +397,7 @@ Please mind that the popular `openai/whisper-large-v3-turbo` model does not supp
|
||||
!!! note
|
||||
To use the Translation API, please install with extra audio dependencies using `pip install vllm[audio]`.
|
||||
|
||||
Code example: [examples/online_serving/openai_translation_client.py](../../examples/online_serving/openai_translation_client.py)
|
||||
Code example: [examples/speech_to_text/openai/openai_translation_client.py](../../examples/speech_to_text/openai/openai_translation_client.py)
|
||||
|
||||
#### Extra Parameters
|
||||
|
||||
|
||||
@@ -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
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user