Compare commits

..
Author SHA1 Message Date
yewentao256 b547993038 type ignore
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-04-21 15:06:50 -04:00
yewentao256 9026d5ec05 assert not None
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-04-21 14:59:44 -04:00
yewentao256 ee778e2cd8 update batch invariant docstring
Signed-off-by: yewentao256 <zhyanwentao@126.com>
2026-04-21 14:47:35 -04:00
1035 changed files with 16455 additions and 86335 deletions
+3 -14
View File
@@ -69,11 +69,11 @@ steps:
pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs
pytest -x -v -s tests/quantization/test_cpu_wna16.py"
- label: CPU-Distributed Tests (PP+TP)
- label: CPU-Distributed Tests
depends_on: []
device: intel_cpu
no_plugin: true
source_file_dependencies: &cpu_distributed_deps
source_file_dependencies:
- csrc/cpu/shm.cpp
- vllm/v1/worker/cpu_worker.py
- vllm/v1/worker/gpu_worker.py
@@ -82,21 +82,10 @@ steps:
- vllm/platforms/cpu.py
- vllm/distributed/parallel_state.py
- vllm/distributed/device_communicators/cpu_communicator.py
- .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 10m "
bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh tp_pp"
- label: CPU-Distributed Tests (DP+TP)
depends_on: []
device: intel_cpu
no_plugin: true
source_file_dependencies: *cpu_distributed_deps
commands:
- |
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 10m "
bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh dp_tp"
bash .buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh"
- label: CPU-Multi-Modal Model Tests %N
depends_on: []
-1
View File
@@ -192,7 +192,6 @@ export BUILDKITE_COMMIT
export PARENT_COMMIT
export IMAGE_TAG
export IMAGE_TAG_LATEST
export COMMIT="${COMMIT:-${BUILDKITE_COMMIT}}"
export CACHE_FROM
export CACHE_FROM_BASE_BRANCH
export CACHE_FROM_MAIN
@@ -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.2"
NIGHTLY_CUDA_VERSION="13.0.0"
NIGHTLY_BUILD_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-devel-ubuntu22.04"
NIGHTLY_FINAL_BASE_IMAGE="nvidia/cuda:${NIGHTLY_CUDA_VERSION}-base-ubuntu22.04"
+3 -2
View File
@@ -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 --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_functions.py &&
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,5 +125,6 @@ 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 || true) &&
pytest -v -s lora/test_qwen3_unembed.py &&
pytest -v -s lora/test_qwenvl.py &&
pytest -v -s lora/test_whisper.py'
+63 -205
View File
@@ -1,13 +1,3 @@
# 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
@@ -24,7 +14,9 @@ steps:
agents:
queue: arm64_cpu_queue_release
commands:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=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 ."
# #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 ."
- "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"
@@ -37,10 +29,12 @@ steps:
agents:
queue: arm64_cpu_queue_release
commands:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_AARCH64}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinuxaarch64-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
# #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 ."
- "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_35"
env:
DOCKER_BUILDKIT: "1"
@@ -53,7 +47,7 @@ steps:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_BUILD_ACL=ON --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
env:
DOCKER_BUILDKIT: "1"
@@ -63,10 +57,10 @@ 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 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86_CU129}\" --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 --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"
@@ -76,10 +70,10 @@ steps:
agents:
queue: cpu_queue_release
commands:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=13.0.2 --build-arg torch_cuda_arch_list=\"${CUDA_ARCH_X86}\" --build-arg BUILD_OS=manylinux --build-arg BUILD_BASE_IMAGE=pytorch/manylinux2_28-builder:cuda13.0 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
- "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 ."
- "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_35"
env:
DOCKER_BUILDKIT: "1"
@@ -92,7 +86,7 @@ steps:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --build-arg VLLM_CPU_X86=true --tag vllm-ci:build-image --target vllm-build --progress plain -f docker/Dockerfile.cpu ."
- "mkdir artifacts"
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_35"
env:
DOCKER_BUILDKIT: "1"
@@ -114,202 +108,102 @@ steps:
depends_on: block-build-release-images
allow_dependency_failure: true
steps:
- 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
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh) \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=13.0.2 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
- "DOCKER_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 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 13.0"
- label: "Build release image - aarch64 - CUDA 12.9"
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh) \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=13.0.2 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu22.04 \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
- "DOCKER_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 push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)"
- 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-cuda-12-9
id: build-release-image-x86-cuda-13-0
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=12.9.1 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
- "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"
# re-tag to default image tag and push, just in case arm64 build fails
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129"
- "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"
- label: "Build release image - aarch64 - CUDA 12.9"
- label: "Build release image - aarch64 - CUDA 13.0"
depends_on: ~
id: build-release-image-arm64-cuda-12-9
id: build-release-image-arm64-cuda-13-0
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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129) \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=12.9.1 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129"
# 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"
- 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-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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=13.0.2 \
--build-arg UBUNTU_VERSION=24.04 \
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
- "DOCKER_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 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 13.0 - Ubuntu 24.04"
- label: "Build release image - aarch64 - CUDA 12.9 - 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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh ubuntu2404) \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=13.0.2 \
--build-arg UBUNTU_VERSION=24.04 \
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--build-arg BUILD_BASE_IMAGE=nvidia/cuda:13.0.2-devel-ubuntu24.04 \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
- "DOCKER_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 push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-ubuntu2404"
- label: "Build release image - x86_64 - CUDA 12.9 - Ubuntu 24.04"
- label: "Build release image - x86_64 - CUDA 13.0 - Ubuntu 24.04"
depends_on: ~
id: build-release-image-x86-cuda-12-9-ubuntu2404
id: build-release-image-x86-cuda-13-0-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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=12.9.1 \
--build-arg UBUNTU_VERSION=24.04 \
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_X86_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
- "docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404 public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-cu129-ubuntu2404"
- "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"
- 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-cuda-12-9-ubuntu2404
id: build-release-image-arm64-cuda-13-0-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 \
$(bash .buildkite/scripts/docker-build-metadata-args.sh cu129-ubuntu2404) \
--build-arg max_jobs=16 \
--build-arg USE_SCCACHE=1 \
--build-arg GIT_REPO_CHECK=1 \
--build-arg CUDA_VERSION=12.9.1 \
--build-arg UBUNTU_VERSION=24.04 \
--build-arg GDRCOPY_OS_VERSION=Ubuntu24_04 \
--build-arg torch_cuda_arch_list="${CUDA_ARCH_AARCH64_CU129}" \
--build-arg INSTALL_KV_CONNECTORS=true \
--target vllm-openai \
--progress plain \
-f docker/Dockerfile .
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT-$(uname -m)-cu129-ubuntu2404"
- "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"
- block: "Build release image for x86_64 CPU"
key: block-cpu-release-image-build
depends_on: ~
- label: "Build release image - x86_64 - CPU"
key: build-cpu-release-image-x86
depends_on:
- block-cpu-release-image-build
- input-release-version
@@ -328,8 +222,7 @@ steps:
depends_on: ~
- label: "Build release image - arm64 - CPU"
key: build-cpu-release-image-arm64
depends_on:
depends_on:
- block-arm64-cpu-release-image-build
- input-release-version
agents:
@@ -345,7 +238,7 @@ steps:
- group: "Publish release images"
key: "publish-release-images"
steps:
- label: "Create multi-arch manifest - CUDA 13.0"
- label: "Create multi-arch manifest - CUDA 12.9"
depends_on:
- build-release-image-x86
- build-release-image-arm64
@@ -357,7 +250,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 13.0"
- label: "Annotate release workflow - CUDA 12.9"
depends_on:
- create-multi-arch-manifest
id: annotate-release-workflow
@@ -366,19 +259,19 @@ steps:
commands:
- "bash .buildkite/scripts/annotate-release.sh"
- label: "Create multi-arch manifest - CUDA 12.9"
- label: "Create multi-arch manifest - CUDA 13.0"
depends_on:
- build-release-image-x86-cuda-12-9
- build-release-image-arm64-cuda-12-9
id: create-multi-arch-manifest-cuda-12-9
- build-release-image-x86-cuda-13-0
- build-release-image-arm64-cuda-13-0
id: create-multi-arch-manifest-cuda-13-0
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-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"
- "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"
- 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-ubuntu2404
- build-release-image-arm64-ubuntu2404
@@ -390,17 +283,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 12.9 - Ubuntu 24.04"
- label: "Create multi-arch manifest - CUDA 13.0 - Ubuntu 24.04"
depends_on:
- 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
- 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
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-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"
- "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"
- label: "Publish nightly multi-arch image to DockerHub"
depends_on:
@@ -420,51 +313,16 @@ steps:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- label: "Publish nightly multi-arch image to DockerHub - CUDA 12.9"
- label: "Publish nightly multi-arch image to DockerHub - CUDA 13.0"
depends_on:
- create-multi-arch-manifest-cuda-12-9
- create-multi-arch-manifest-cuda-13-0
if: build.env("NIGHTLY") == "1"
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/push-nightly-builds.sh cu129"
- "bash .buildkite/scripts/push-nightly-builds.sh cu130"
# Clean up old nightly builds (keep only last 14)
- "bash .buildkite/scripts/cleanup-nightly-builds.sh cu129-nightly-"
plugins:
- docker-login#v3.0.0:
username: vllmbot
password-env: DOCKERHUB_TOKEN
env:
DOCKER_BUILDKIT: "1"
DOCKERHUB_USERNAME: "vllmbot"
- block: "Publish release images to DockerHub"
key: block-publish-release-images
depends_on:
- create-multi-arch-manifest
- create-multi-arch-manifest-cuda-12-9
- create-multi-arch-manifest-ubuntu2404
- create-multi-arch-manifest-cuda-12-9-ubuntu2404
- build-rocm-release-image
- input-release-version
# Wait for CPU builds if their block steps were unblocked, so publish
# doesn't race the in-progress CPU build. allow_failure lets publish
# proceed when the operator legitimately leaves the CPU block steps
# unblocked or the CPU build fails.
- step: build-cpu-release-image-x86
allow_failure: true
- step: build-cpu-release-image-arm64
allow_failure: true
if: build.env("NIGHTLY") != "1"
- label: "Publish release images to DockerHub"
depends_on:
- block-publish-release-images
key: publish-release-images-dockerhub
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/publish-release-images.sh"
- "bash .buildkite/scripts/cleanup-nightly-builds.sh cu130-nightly-"
plugins:
- docker-login#v3.0.0:
username: vllmbot
@@ -760,7 +618,7 @@ steps:
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
env:
S3_BUCKET: "vllm-wheels"
VARIANT: "rocm722"
VARIANT: "rocm721"
# ROCm Job 6: Build ROCm Release Docker Image
- label: ":docker: Build release image - x86_64 - ROCm"
+98 -6
View File
@@ -8,20 +8,112 @@ if [ -z "${RELEASE_VERSION}" ]; then
RELEASE_VERSION="1.0.0.dev"
fi
ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
buildkite-agent annotate --style 'info' --context 'release-workflow' << EOF
To download the wheel (by commit):
\`\`\`
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 .
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 .
(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 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 CPU:
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_x86_64.whl .
aws s3 cp s3://vllm-wheels/${BUILDKITE_COMMIT}/vllm-${RELEASE_VERSION}+cpu-cp38-abi3-manylinux_2_35_aarch64.whl .
\`\`\`
Docker images are published automatically by the "Publish release images to DockerHub" pipeline step.
To download and upload the image:
\`\`\`
# Download images:
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-cu130
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-cu130
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm
docker pull public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
docker pull public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
# Tag and push images:
## CUDA
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64 vllm/vllm-openai:x86_64
docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:latest-x86_64
docker tag vllm/vllm-openai:x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
docker push vllm/vllm-openai:latest-x86_64
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-x86_64-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}-aarch64 vllm/vllm-openai:aarch64
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:latest-aarch64
docker tag vllm/vllm-openai:aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker push vllm/vllm-openai:latest-aarch64
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-aarch64-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
## ROCm
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${BUILDKITE_COMMIT}-rocm vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:latest
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT} vllm/vllm-openai-rocm:v${RELEASE_VERSION}
docker push vllm/vllm-openai-rocm:latest
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:latest-base
docker tag vllm/vllm-openai-rocm:${BUILDKITE_COMMIT}-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
docker push vllm/vllm-openai-rocm:latest-base
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
## CPU
docker tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:x86_64
docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:latest-x86_64
docker tag vllm/vllm-openai-cpu:x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
docker push vllm/vllm-openai-cpu:latest-x86_64
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
docker tag public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} vllm/vllm-openai-cpu:arm64
docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:latest-arm64
docker tag vllm/vllm-openai-cpu:arm64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
docker push vllm/vllm-openai-cpu:latest-arm64
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
# Create multi-arch manifest:
docker manifest rm vllm/vllm-openai:latest
docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker manifest push vllm/vllm-openai:latest
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
docker manifest rm vllm/vllm-openai:latest-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-cpu:latest || true
docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
docker manifest push vllm/vllm-openai-cpu:latest
docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
\`\`\`
EOF
@@ -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/cu130"
TORCH_INDEX_URL="https://download.pytorch.org/whl/cu129"
fi
echo ">>> Using PyTorch index: ${TORCH_INDEX_URL:-PyPI default}"
-55
View File
@@ -1,55 +0,0 @@
#!/bin/bash
# Usage: ./ci-fetch-log.sh <buildkite_job_url> [output_file]
# ./ci-fetch-log.sh <build_number> <job_uuid> [output_file]
#
# Downloads the raw log for a Buildkite job from the public, unauthenticated
# /organizations/<org>/pipelines/<pipeline>/builds/<n>/jobs/<uuid>/download
# endpoint, then strips ANSI/timestamps via ci-clean-log.sh.
#
# Find <build_number> and <job_uuid> via:
# gh pr checks <PR> --repo vllm-project/vllm
# Each failing row's URL is .../builds/<build_number>#<job_uuid>.
set -euo pipefail
ORG="vllm"
PIPELINE="ci"
usage() {
echo "Usage: $0 <buildkite_job_url> [output_file]"
echo " $0 <build_number> <job_uuid> [output_file]"
exit 1
}
if [ $# -lt 1 ]; then usage; fi
if [[ "$1" == https://* ]]; then
BUILD=$(echo "$1" | sed -nE 's#.*/builds/([0-9]+).*#\1#p')
JOB=$(echo "$1" | grep -oE '[0-9a-f]{8}-[0-9a-f-]+' | head -n 1)
OUT="${2:-ci-${BUILD}-${JOB:0:8}.log}"
else
if [ $# -lt 2 ]; then usage; fi
BUILD="$1"
JOB="$2"
OUT="${3:-ci-${BUILD}-${JOB:0:8}.log}"
fi
if [ -z "$BUILD" ] || [ -z "$JOB" ]; then
echo "Could not parse build number or job UUID from: $1" >&2
usage
fi
COOKIES=$(mktemp)
trap 'rm -f "$COOKIES"' EXIT
# Buildkite issues a session cookie on first hit; subsequent /download needs it.
curl -fsSL -c "$COOKIES" -A "vllm-ci-fetch-log" \
"https://buildkite.com/${ORG}/${PIPELINE}/builds/${BUILD}" -o /dev/null
curl -fsSL -b "$COOKIES" -A "vllm-ci-fetch-log" \
"https://buildkite.com/organizations/${ORG}/pipelines/${PIPELINE}/builds/${BUILD}/jobs/${JOB}/download" \
-o "$OUT"
bash "$(dirname "$0")/ci-clean-log.sh" "$OUT"
echo "$OUT"
-142
View File
@@ -1,142 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Detect the manylinux platform tag for a wheel and rename it in place.
vLLM's build images produce wheels with the generic ``linux_<arch>`` platform
tag, which installers like ``pip`` won't accept off PyPI/our index. We need to
rewrite the platform tag to the appropriate ``manylinux_<major>_<minor>_<arch>``
before uploading.
Historically the tag was hard-coded per build (``manylinux_2_31`` for the
Ubuntu 20.04-based image, ``manylinux_2_35`` for the Ubuntu 22.04-based
images). That is brittle: bumping the base image silently produces wheels
labelled with the wrong glibc requirement. This script asks ``auditwheel``
to derive the tag from the symbol versions actually referenced by the
binaries inside the wheel, so the label tracks reality.
We can't simply call ``auditwheel repair`` -- it tries to graft external
shared libraries into the wheel and fails on vLLM's CUDA/cuBLAS dependencies.
Instead we use ``auditwheel.wheel_abi.analyze_wheel_abi`` directly, which is
the same call that powers ``auditwheel show``, and read off
``winfo.sym_policy.name``.
Usage:
detect-manylinux-tag.py <wheel_path>
The wheel is renamed in place; the new path is printed on stdout. All
diagnostics go to stderr so callers can capture stdout safely.
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
from auditwheel.error import (
AuditwheelError,
NonPlatformWheelError,
WheelToolsError,
)
from auditwheel.wheel_abi import analyze_wheel_abi
from auditwheel.wheeltools import get_wheel_architecture, get_wheel_libc
def detect_platform_tag(wheel_path: Path) -> str:
"""Return the most precise platform tag the wheel is consistent with.
Mirrors ``auditwheel show`` but returns ``sym_policy`` rather than
``overall_policy``: we only care about the glibc symbol versions used,
not about other policy axes (ISA extensions, blacklist, etc.) that
``overall_policy`` folds in.
"""
fn = wheel_path.name
try:
arch = get_wheel_architecture(fn)
except (WheelToolsError, NonPlatformWheelError):
# Architecture isn't deducible from the filename; let auditwheel
# infer it from the ELF binaries inside the wheel.
arch = None
try:
libc = get_wheel_libc(fn)
except WheelToolsError:
# An unrepaired wheel uses ``linux_<arch>``, which doesn't encode
# libc. Let auditwheel infer it from the ELF binaries.
libc = None
winfo = analyze_wheel_abi(
libc,
arch,
wheel_path,
frozenset(),
disable_isa_ext_check=False,
allow_graft=False,
)
return winfo.sym_policy.name
def rename_wheel(wheel_path: Path, new_platform_tag: str) -> Path:
"""Rename the wheel in place, replacing only its platform tag."""
# Wheel filename per PEP 427:
# {distribution}-{version}(-{build})?-{python}-{abi}-{platform}.whl
# The platform tag is always the last ``-``-separated token before
# ``.whl``. Compound tags like ``manylinux_2_31_x86_64`` use ``_`` as the
# internal separator, so ``-``-splitting is unambiguous.
parts = wheel_path.stem.split("-")
if len(parts) < 5:
raise ValueError(f"Unrecognised wheel filename: {wheel_path.name}")
parts[-1] = new_platform_tag
new_path = wheel_path.with_name("-".join(parts) + ".whl")
if new_path != wheel_path:
wheel_path.rename(new_path)
return new_path
def main() -> int:
parser = argparse.ArgumentParser(
description="Detect a wheel's manylinux platform tag with "
"auditwheel and rename the wheel in place."
)
parser.add_argument(
"wheel",
type=Path,
help="Path to the wheel to inspect and rename.",
)
args = parser.parse_args()
wheel_path: Path = args.wheel
if not wheel_path.is_file():
print(f"error: {wheel_path} is not a file", file=sys.stderr)
return 1
# Catch the things that ``analyze_wheel_abi`` and ``rename_wheel`` can
# raise: any subclass of ``AuditwheelError`` (pure-Python wheels,
# invalid libc, malformed wheels), filesystem errors, or our own
# ``ValueError`` for an unrecognised wheel filename. Print a single
# ``ERROR_TYPE: message`` line to stderr instead of a Python
# traceback, which is much friendlier in CI logs.
try:
new_tag = detect_platform_tag(wheel_path)
print(f"detected platform tag: {new_tag}", file=sys.stderr)
new_path = rename_wheel(wheel_path, new_tag)
except (AuditwheelError, ValueError, OSError) as e:
print(
f"error: failed to retag {wheel_path.name}: {type(e).__name__}: {e}",
file=sys.stderr,
)
return 2
if new_path != wheel_path:
print(f"renamed {wheel_path.name} -> {new_path.name}", file=sys.stderr)
else:
print(f"wheel already tagged {new_tag}", file=sys.stderr)
print(new_path)
return 0
if __name__ == "__main__":
sys.exit(main())
@@ -1,54 +0,0 @@
#!/bin/bash
# Emit docker build flags for release image provenance metadata.
# Keep this helper best-effort: missing Buildkite metadata should fall back to
# local/default values instead of blocking the Docker build.
# Variant examples: "", "cu129", "ubuntu2404", "cu129-ubuntu2404".
variant="${1:-}"
variant_suffix="${variant:+-${variant}}"
image_name="${VLLM_DOCKER_IMAGE_NAME:-vllm/vllm-openai}"
staging_repo="${VLLM_STAGING_IMAGE_REPO:-public.ecr.aws/q9t5s3a7/vllm-release-repo}"
build_commit="${VLLM_BUILD_COMMIT:-${BUILDKITE_COMMIT:-unknown}}"
build_pipeline="${VLLM_BUILD_PIPELINE:-${BUILDKITE_PIPELINE_ID:-${BUILDKITE_PIPELINE_SLUG:-local}}}"
build_url="${VLLM_BUILD_URL:-${BUILDKITE_BUILD_URL:-}}"
tag_commit="${BUILDKITE_COMMIT:-${build_commit}}"
if [[ -n "${BUILDKITE:-}" || -n "${BUILDKITE_COMMIT:-}" ]]; then
release_version="${RELEASE_VERSION:-}"
if command -v buildkite-agent >/dev/null 2>&1; then
release_version="${release_version:-$(buildkite-agent meta-data get release-version 2>/dev/null)}"
fi
release_version="${release_version#v}"
release_version="${release_version:-${tag_commit}}"
staging_image_ref="${staging_repo}:${tag_commit}-$(uname -m)${variant_suffix}"
if [[ "${NIGHTLY:-}" == "1" ]]; then
if [[ -z "${variant}" ]]; then
image_tag="${image_name}:nightly-${tag_commit}"
elif [[ "${variant}" == cu* ]]; then
cuda_variant="${variant%%-*}"
remaining_variant="${variant#${cuda_variant}}"
image_tag="${image_name}:${cuda_variant}-nightly-${tag_commit}${remaining_variant}"
else
image_tag="${image_name}:nightly-${tag_commit}${variant_suffix}"
fi
else
image_tag="${image_name}:v${release_version}${variant_suffix}"
fi
else
image_tag="${VLLM_IMAGE_TAG:-local/vllm-openai:dev}"
staging_image_ref="${image_tag}"
fi
emit_arg() {
printf -- "--build-arg %s=%s " "$1" "$2"
}
emit_arg VLLM_BUILD_COMMIT "${build_commit}"
emit_arg VLLM_BUILD_PIPELINE "${build_pipeline}"
emit_arg VLLM_BUILD_URL "${build_url}"
# This is the intended public tag. The final digest is only known after push.
emit_arg VLLM_IMAGE_TAG "${image_tag}"
printf -- "--tag %s " "${staging_image_ref}"
@@ -9,14 +9,21 @@ set -ex
BUCKET="vllm-wheels"
INDICES_OUTPUT_DIR="indices"
DEFAULT_VARIANT_ALIAS="cu130" # align with vLLM_MAIN_CUDA_VERSION in vllm/envs.py
DEFAULT_VARIANT_ALIAS="cu129" # 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/"
# Select python3 (>= 3.12) -- local if available, else a docker fallback.
# shellcheck source=lib/select-python.sh
source .buildkite/scripts/lib/select-python.sh
select_python
# detect if python3.12+ is available
has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
if [[ "$has_new_python" -eq 0 ]]; then
# use new python from docker
docker pull python:3-slim
PYTHON="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
fi
echo "Using python interpreter: $PYTHON"
echo "Python version: $($PYTHON --version)"
# ======== generate and upload indices ========
@@ -3,37 +3,42 @@ set -euox pipefail
export VLLM_CPU_CI_ENV=0
export VLLM_CPU_KVCACHE_SPACE=1 # avoid OOM
MODE=${1:-all}
echo "--- PP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -pp=2 --max-model-len=4096 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model meta-llama/Llama-3.2-3B-Instruct \
--num-prompts 20 \
--result-dir ./test_results \
--result-filename tp_pp.json \
--save-result \
--endpoint /v1/completions
kill -s SIGTERM $server_pid; wait $server_pid || true
failed_req=$(jq '.failed' ./test_results/tp_pp.json)
if [ "$failed_req" -ne 0 ]; then
echo "Some requests were failed!"
exit 1
fi
run_scenario() {
local label="$1" result_file="$2"
shift 2
echo "--- $label"
vllm serve meta-llama/Llama-3.2-3B-Instruct "$@" --max-model-len=4096 &
local server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model meta-llama/Llama-3.2-3B-Instruct \
--num-prompts 20 \
--result-dir ./test_results \
--result-filename "$result_file" \
--save-result \
--endpoint /v1/completions
kill -s SIGTERM "$server_pid"; wait "$server_pid" || true
if [ "$(jq '.failed' "./test_results/$result_file")" -ne 0 ]; then
echo "Some requests were failed in $label!"
exit 1
fi
}
case "$MODE" in
tp_pp) run_scenario "PP+TP" tp_pp.json -tp=2 -pp=2 ;;
dp_tp) run_scenario "DP+TP" dp_tp.json -tp=2 -dp=2 ;;
all)
run_scenario "PP+TP" tp_pp.json -tp=2 -pp=2
run_scenario "DP+TP" dp_tp.json -tp=2 -dp=2
;;
*) echo "ERROR: unknown mode '$MODE' (expected: tp_pp | dp_tp | all)" >&2; exit 1 ;;
esac
echo "--- DP+TP"
vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
server_pid=$!
timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
vllm bench serve \
--backend vllm \
--dataset-name random \
--model meta-llama/Llama-3.2-3B-Instruct \
--num-prompts 20 \
--result-dir ./test_results \
--result-filename dp_pp.json \
--save-result \
--endpoint /v1/completions
kill -s SIGTERM $server_pid; wait $server_pid || true
failed_req=$(jq '.failed' ./test_results/dp_pp.json)
if [ "$failed_req" -ne 0 ]; then
echo "Some requests were failed!"
exit 1
fi
@@ -25,100 +25,22 @@ 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
flock -u 9
return 1
exit 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%. Running aggressive CI image cleanup..."
cleanup_old_ci_images "${REGISTRY}/${REPO}" "${image_name}" "${DOCKER_IMAGE_CLEANUP_HOURS:-72}" 1
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."
else
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)."
echo "Disk usage is below $threshold%. No cleanup needed."
fi
}
@@ -343,6 +265,8 @@ 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,7 +12,9 @@ docker build -t "${image_name}" -f docker/Dockerfile.xpu .
# Setup cleanup
remove_docker_container() {
docker rm -f "${container_name}" || true
docker rm -f "${container_name}" || true;
docker image rm -f "${image_name}" || true;
docker system prune -f || true;
}
trap remove_docker_container EXIT
-127
View File
@@ -1,127 +0,0 @@
#!/usr/bin/env bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Shared helper for rewriting a wheel's platform tag from the generic
# ``linux_<arch>`` to the correct ``manylinux_<major>_<minor>_<arch>``.
# After sourcing, call ``apply_manylinux_tag <wheel>`` on each wheel
# that still carries the generic tag; the renamed path is printed on
# stdout (logs go to stderr).
#
# Why a pinned Docker container instead of using whatever Python
# happens to be on the agent:
# - vLLM's release agents are heterogeneous -- they don't agree on
# a Python minor version, and we can't rely on a particular
# ``auditwheel`` being installed.
# - ``detect-manylinux-tag.py`` reads ``auditwheel.wheel_abi`` and
# ``Policy.sym_policy``, which are *internal* APIs without a
# stability promise. Pinning both Python and auditwheel makes the
# detected tag a function of the inputs alone, and shifts version
# bumps from "implicit drift" to "deliberate, retested change".
# - Other release scripts (``generate-and-upload-nightly-index.sh``,
# ``upload-rocm-wheels.sh``) already use the python:3-slim image
# when the agent's interpreter is too old; this is the same idea
# made stricter.
#
# To keep the per-wheel cost down (the ROCm upload retags ~10 wheels
# each run), we install auditwheel into a long-lived helper container
# once on source, then ``docker exec`` into it for each call.
#
# Trap behaviour:
# - Sourcing installs an EXIT trap that calls ``manylinux_cleanup`` to
# tear down the helper container. Any EXIT trap that was already in
# place when this file was sourced is captured and run AFTER our
# cleanup, so we don't silently clobber it.
# - If a caller sets a new EXIT trap *after* sourcing, that trap will
# replace ours; in that case the caller should call
# ``manylinux_cleanup`` from their own handler.
if [[ -n "${_MANYLINUX_LIB_SOURCED:-}" ]]; then
return 0
fi
_MANYLINUX_LIB_SOURCED=1
# Pin both sides. Bump these deliberately and re-run a representative
# wheel from each build target through the detection.
_MANYLINUX_PYTHON_IMAGE="python:3.12-slim"
_MANYLINUX_AUDITWHEEL_VERSION="6.6.0"
# Resolve our own directory (and the sibling detect script) using the
# canonical, symlink-resolved path. The container mounts cwd at the
# same absolute path on both sides, so all paths we hand to it -- the
# script, the wheel -- must canonicalise to a location under cwd.
_MANYLINUX_LIB_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)"
_MANYLINUX_DETECT_SCRIPT="$(cd "${_MANYLINUX_LIB_DIR}/.." && pwd -P)/detect-manylinux-tag.py"
_MANYLINUX_CWD="$(pwd -P)"
docker pull --quiet "$_MANYLINUX_PYTHON_IMAGE" >/dev/null
# Spin up a long-lived helper container so we install auditwheel once
# and then ``docker exec`` into it for each wheel.
#
# The container runs as root so ``pip install`` can write into the
# system site-packages; individual ``docker exec`` calls below pin
# themselves to the host UID so any file rename happens with host
# ownership, not root.
_MANYLINUX_CONTAINER="$(docker run -d --rm \
-v "$_MANYLINUX_CWD:$_MANYLINUX_CWD" \
-w "$_MANYLINUX_CWD" \
"$_MANYLINUX_PYTHON_IMAGE" \
sleep infinity)"
docker exec "$_MANYLINUX_CONTAINER" \
pip install --quiet --disable-pip-version-check \
--root-user-action=ignore \
"auditwheel==${_MANYLINUX_AUDITWHEEL_VERSION}"
# Public cleanup -- safe to call multiple times.
manylinux_cleanup() {
if [[ -n "${_MANYLINUX_CONTAINER:-}" ]]; then
docker rm -f "$_MANYLINUX_CONTAINER" >/dev/null 2>&1 || true
_MANYLINUX_CONTAINER=""
fi
}
# Capture any EXIT trap that was already in place so we can chain to
# it rather than overwrite it. ``trap -p EXIT`` prints the handler in
# eval-able form (``trap -- 'CMD' EXIT``) or nothing if unset; we
# strip the wrapper to recover ``CMD``. Handles the common case --
# CMDs without embedded single quotes -- and degrades gracefully (we
# still run our own cleanup) for the pathological case.
_manylinux_prev_exit_trap_cmd=""
_manylinux_existing_exit_trap="$(trap -p EXIT)"
if [[ -n "$_manylinux_existing_exit_trap" ]]; then
_tmp="${_manylinux_existing_exit_trap#trap -- \'}"
_manylinux_prev_exit_trap_cmd="${_tmp%\' EXIT}"
unset _tmp
fi
unset _manylinux_existing_exit_trap
_manylinux_run_exit_chain() {
manylinux_cleanup
if [[ -n "$_manylinux_prev_exit_trap_cmd" ]]; then
eval "$_manylinux_prev_exit_trap_cmd"
fi
}
trap _manylinux_run_exit_chain EXIT
# Detect the manylinux platform tag for a single wheel and rename it
# in place, printing the renamed wheel path on stdout. Returns
# non-zero on failure (which under ``set -e`` propagates to caller).
#
# The wheel must be reachable via a path under the host cwd so it's
# visible inside the helper container; in CI the wheels always live
# under ``artifacts/`` so this is fine.
apply_manylinux_tag() {
local wheel="$1"
local abs_wheel
abs_wheel="$(realpath "$wheel")"
local new_wheel
new_wheel="$(docker exec -u "$(id -u):$(id -g)" \
"$_MANYLINUX_CONTAINER" \
python "$_MANYLINUX_DETECT_SCRIPT" "$abs_wheel")"
if [[ -z "$new_wheel" || ! -f "$new_wheel" ]]; then
echo "apply_manylinux_tag: detect-manylinux-tag.py did not produce a valid wheel path for $wheel" >&2
return 1
fi
printf '%s\n' "$new_wheel"
}
-41
View File
@@ -1,41 +0,0 @@
#!/usr/bin/env bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Pick a Python interpreter for buildkite scripts: prefer a local
# ``python3`` if it is recent enough (>= 3.12), otherwise fall back to
# a one-shot Docker container running ``python:3-slim``. After
# ``select_python`` returns, ``$PYTHON`` is set in the caller's shell
# and is safe to use as a command (e.g. ``$PYTHON some_script.py``).
#
# The 3.12 threshold matches what the existing nightly-index work
# expects -- typing features used by ``generate-nightly-index.py``.
# This helper does not pin the *minor* version; if you need stricter
# reproducibility (e.g. relying on auditwheel internals), invoke
# Docker yourself with a pinned tag rather than calling this.
if [[ -n "${_SELECT_PYTHON_LIB_SOURCED:-}" ]]; then
return 0
fi
_SELECT_PYTHON_LIB_SOURCED=1
# Sets ``PYTHON`` in the caller's shell and exports it. Idempotent --
# calling twice is safe and the second call simply re-runs the probe.
select_python() {
local py="${PYTHON_PROG:-python3}"
local has_new_python
has_new_python=$("$py" -c \
"print(1 if __import__('sys').version_info >= (3,12) else 0)" \
2>/dev/null || echo 0)
if [[ "$has_new_python" -eq 0 ]]; then
# ``-u $(id -u):$(id -g)`` so files created via the container
# end up owned by the host user, not root.
docker pull python:3-slim
PYTHON="docker run --rm -u $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
else
PYTHON="$py"
fi
export PYTHON
echo "Using python interpreter: $PYTHON"
echo "Python version: $($PYTHON --version)"
}
@@ -1,180 +0,0 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Publish release Docker images from ECR to DockerHub.
# Pulls per-arch images, tags with latest and versioned tags, pushes them,
# then creates and pushes multi-arch manifests.
set -euo pipefail
RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//')
if [ -z "${RELEASE_VERSION}" ]; then
echo "ERROR: release-version metadata not set"
exit 1
fi
COMMIT="$BUILDKITE_COMMIT"
ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
echo "========================================"
echo "Publishing release images v${RELEASE_VERSION}"
echo " Commit: ${COMMIT}"
echo " ROCm base cache key: ${ROCM_BASE_CACHE_KEY}"
echo "========================================"
# Login to ECR to pull staging images
aws ecr-public get-login-password --region us-east-1 | \
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
# ---- CUDA (default: 13.0) ----
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
docker push vllm/vllm-openai:latest-x86_64
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker push vllm/vllm-openai:latest-aarch64
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker manifest rm vllm/vllm-openai:latest || true
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true
docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64
docker manifest push vllm/vllm-openai:latest
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}
# ---- CUDA 12.9 ----
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
docker push vllm/vllm-openai:latest-x86_64-cu129
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
docker push vllm/vllm-openai:latest-aarch64-cu129
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
docker manifest rm vllm/vllm-openai:latest-cu129 || true
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true
docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129
docker manifest push vllm/vllm-openai:latest-cu129
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129
# ---- Ubuntu 24.04 (CUDA 13.0) ----
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404
docker push vllm/vllm-openai:latest-x86_64-ubuntu2404
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
docker push vllm/vllm-openai:latest-aarch64-ubuntu2404
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true
docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404
docker manifest push vllm/vllm-openai:latest-ubuntu2404
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404
# ---- Ubuntu 24.04 (CUDA 12.9) ----
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404
docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404
docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true
docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true
docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404
docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404
docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404
docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404
# ---- ROCm ----
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION}
docker push vllm/vllm-openai-rocm:latest
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
docker push vllm/vllm-openai-rocm:latest-base
docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base
# ---- CPU ----
# CPU images are behind separate block steps and may not have been built.
# All-or-nothing: inspect both arches first, then either publish everything
# (per-arch + multi-arch manifest) or skip everything. Publishing only one
# arch would leave `:latest-x86_64` pointing at the new release while the
# `:latest` multi-arch manifest still resolves to the previous release.
CPU_X86_TAG=public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION}
CPU_ARM_TAG=public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION}
CPU_X86_AVAILABLE=false
CPU_ARM_AVAILABLE=false
docker manifest inspect "${CPU_X86_TAG}" >/dev/null 2>&1 && CPU_X86_AVAILABLE=true
docker manifest inspect "${CPU_ARM_TAG}" >/dev/null 2>&1 && CPU_ARM_AVAILABLE=true
if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then
docker pull "${CPU_X86_TAG}"
docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64
docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
docker push vllm/vllm-openai-cpu:latest-x86_64
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64
docker pull "${CPU_ARM_TAG}"
docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64
docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
docker push vllm/vllm-openai-cpu:latest-arm64
docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
docker manifest rm vllm/vllm-openai-cpu:latest || true
docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true
docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64
docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64
docker manifest push vllm/vllm-openai-cpu:latest
docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION}
elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then
echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)"
else
# Partial state: one arch built, the other did not. Fail loudly rather than
# ship a Docker Hub state where `:latest-${arch}` and `:latest` (multi-arch)
# disagree on which release they point at.
echo "ERROR: Partial CPU build detected (x86_64=${CPU_X86_AVAILABLE}, arm64=${CPU_ARM_AVAILABLE})."
echo " Refusing to publish to avoid split-tag drift between per-arch and multi-arch tags."
echo " Re-run the missing CPU build and retry, or manually publish if a single-arch release is intended."
exit 1
fi
echo ""
echo "Successfully published release images for v${RELEASE_VERSION}"
@@ -51,7 +51,6 @@ vllm serve "$MODEL" \
--offload-num-in-group 2 \
--offload-prefetch-step 1 \
--offload-params w13_weight w2_weight \
--generation-config vllm \
--port "$PORT" \
${EXTRA_ARGS+"${EXTRA_ARGS[@]}"} &
SERVER_PID=$!
@@ -28,7 +28,6 @@
# BFCL_MAX_MODEL_LEN - Max model length (default: 4096)
# BFCL_PORT - Server port (default: 8000)
# BFCL_REASONING_PARSER - Reasoning parser name (default: disabled)
# BFCL_TEMPERATURE - Temperature (default: 0.0)
# BFCL_EXTRA_ARGS - Additional vLLM server args
set -euo pipefail
@@ -44,7 +43,6 @@ TP_SIZE="${BFCL_TP_SIZE:-1}"
MAX_MODEL_LEN="${BFCL_MAX_MODEL_LEN:-4096}"
PORT="${BFCL_PORT:-8000}"
REASONING_PARSER="${BFCL_REASONING_PARSER:-}"
TEMPERATURE="${BFCL_TEMPERATURE:-0.0}"
EXTRA_ARGS="${BFCL_EXTRA_ARGS:-}"
# Set up output directory
@@ -141,7 +139,7 @@ echo "vLLM server is ready. (started in ${SECONDS_WAITED}s)"
# be patched in-process so BFCL knows to use the OpenAI-compatible handler
# against our local vLLM server.
bfcl_exit_code=0
python3 - "$MODEL" "$TEST_CATEGORY" "$NUM_THREADS" "$PORT" "$API_TYPE" "$TEMPERATURE" "$OUTPUT_DIR" << 'PYEOF' || bfcl_exit_code=$?
python3 - "$MODEL" "$TEST_CATEGORY" "$NUM_THREADS" "$PORT" "$API_TYPE" "$OUTPUT_DIR" << 'PYEOF' || bfcl_exit_code=$?
import os
import sys
@@ -150,8 +148,7 @@ test_category = sys.argv[2]
num_threads = int(sys.argv[3])
port = sys.argv[4]
api_type = sys.argv[5]
temperature = float(sys.argv[6])
output_dir = sys.argv[7] if len(sys.argv) > 7 and sys.argv[7] else os.getcwd()
output_dir = sys.argv[6] if len(sys.argv) > 6 and sys.argv[6] else os.getcwd()
os.environ["OPENAI_BASE_URL"] = f"http://localhost:{port}/v1"
os.environ["OPENAI_API_KEY"] = "dummy"
@@ -207,7 +204,6 @@ gen_kwargs["model"] = [model]
gen_kwargs["test_category"] = [c.strip() for c in test_category.split(",")]
gen_kwargs["skip_server_setup"] = True
gen_kwargs["num_threads"] = num_threads
gen_kwargs["temperature"] = temperature
generate(**gen_kwargs)
# ---- evaluate ----
+14 -8
View File
@@ -2,18 +2,14 @@
set -ex
# Upload a single wheel to S3, after detecting and applying the appropriate
# manylinux platform tag with auditwheel.
# Upload a single wheel to S3 (rename linux -> manylinux).
# Index generation is handled separately by generate-and-upload-nightly-index.sh.
# shellcheck source=lib/manylinux.sh
source .buildkite/scripts/lib/manylinux.sh
BUCKET="vllm-wheels"
SUBPATH=$BUILDKITE_COMMIT
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
# ========= locate the wheel ==========
# ========= collect, rename & upload the wheel ==========
# Assume wheels are in artifacts/dist/*.whl
wheel_files=(artifacts/dist/*.whl)
@@ -25,9 +21,19 @@ if [[ ${#wheel_files[@]} -ne 1 ]]; then
fi
wheel="${wheel_files[0]}"
# ========= detect manylinux tag and rename ==========
# default build image uses ubuntu 20.04, which corresponds to manylinux_2_31
# we also accept params as manylinux tag
# refer to https://github.com/mayeut/pep600_compliance?tab=readme-ov-file#acceptable-distros-to-build-wheels
manylinux_version="${1:-manylinux_2_31}"
wheel="$(apply_manylinux_tag "$wheel")"
# Rename 'linux' to the appropriate manylinux version in the wheel filename
if [[ "$wheel" != *"linux"* ]]; then
echo "Error: Wheel filename does not contain 'linux': $wheel"
exit 1
fi
new_wheel="${wheel/linux/$manylinux_version}"
mv -- "$wheel" "$new_wheel"
wheel="$new_wheel"
echo "Renamed wheel to: $wheel"
# Extract the version from the wheel
@@ -39,11 +39,10 @@ fi
set -x # avoid printing secrets above
# install twine and sdist build prerequisites from pypi
# install twine from pypi
python3 -m venv /tmp/vllm-release-env
source /tmp/vllm-release-env/bin/activate
pip install twine
pip install -r requirements/build/cuda.txt
python3 -m twine --version
# copy release wheels to local directory
+18 -23
View File
@@ -20,6 +20,10 @@ BUCKET="${S3_BUCKET:-vllm-wheels}"
ROCM_SUBPATH="rocm/${BUILDKITE_COMMIT}"
S3_COMMIT_PREFIX="s3://$BUCKET/$ROCM_SUBPATH/"
INDICES_OUTPUT_DIR="rocm-indices"
PYTHON="${PYTHON_PROG:-python3}"
# ROCm uses manylinux_2_35 (Ubuntu 22.04 based)
MANYLINUX_VERSION="manylinux_2_35"
echo "========================================"
echo "ROCm Wheel Upload Configuration"
@@ -30,21 +34,19 @@ echo "Commit: $BUILDKITE_COMMIT"
echo "Branch: $BUILDKITE_BRANCH"
echo "========================================"
# ======== Part 0: Setup Python and helpers ========
# ======== Part 0: Setup Python ========
# Pick a Python interpreter for index generation -- local if recent
# enough, else a one-shot docker fallback.
# shellcheck source=lib/select-python.sh
source .buildkite/scripts/lib/select-python.sh
select_python
# Detect if python3.12+ is available
has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)" 2>/dev/null || echo 0)
if [[ "$has_new_python" -eq 0 ]]; then
# Use new python from docker
# Use --user to ensure files are created with correct ownership (not root)
docker pull python:3-slim
PYTHON="docker run --rm --user $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
fi
# Set up auditwheel-in-a-container for the manylinux retagging step.
# Distinct from select_python: ``manylinux.sh`` deliberately pins both
# the Python and auditwheel versions (the script reads auditwheel
# internals) and so always runs in a known-good container regardless
# of what's on the agent.
# shellcheck source=lib/manylinux.sh
source .buildkite/scripts/lib/manylinux.sh
echo "Using python interpreter: $PYTHON"
echo "Python version: $($PYTHON --version)"
# ======== Part 1: Collect and prepare wheels ========
@@ -61,18 +63,11 @@ if [ "$WHEEL_COUNT" -eq 0 ]; then
exit 1
fi
# Detect the appropriate manylinux platform tag for any wheel that still
# carries the generic ``linux_<arch>`` tag, and rename it in place. We use
# auditwheel via ``apply_manylinux_tag`` (see lib/manylinux.sh) rather than
# a hard-coded ``manylinux_2_35`` string so that the label tracks the actual
# glibc symbol versions used by the binaries (and stays correct if the
# rocm_base image is rebased).
#
# The ``linux``/``manylinux`` filter below skips both pre-tagged wheels
# (e.g. upstream torch) and pure-Python ``-any.whl`` wheels.
# Rename linux to manylinux in wheel filenames
for wheel in all-rocm-wheels/*.whl; do
if [[ "$wheel" == *"linux"* ]] && [[ "$wheel" != *"manylinux"* ]]; then
new_wheel="$(apply_manylinux_tag "$wheel")"
new_wheel="${wheel/linux/$MANYLINUX_VERSION}"
mv -- "$wheel" "$new_wheel"
echo "Renamed: $(basename "$wheel") -> $(basename "$new_wheel")"
fi
done
+31 -32
View File
@@ -388,18 +388,18 @@ steps:
- python3 basic/offline_inference/embed.py
- python3 basic/offline_inference/score.py
# Multi-modal models
- 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
- 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
# Pooling models
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# Features demo
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 offline_inference/prefix_caching.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 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
- 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
#---------------------------------------------------------- mi250 · kernels ----------------------------------------------------------#
@@ -1108,7 +1108,6 @@ steps:
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
- VLLM_TEST_CLEAN_GPU_MEMORY=1 pytest -v -s tests/compile/passes/distributed/test_async_tp.py
- pytest -v -s tests/compile/passes/distributed/test_sequence_parallelism.py
- pytest -v -s tests/compile/passes/distributed/test_tp2_ar_rms.py::test_tp2_ar_rms_fusions
#----------------------------------------------------------- mi300 · cuda ------------------------------------------------------------#
@@ -1169,13 +1168,13 @@ steps:
- vllm/v1/attention/backends/
- vllm/v1/attention/selector.py
- tests/distributed/test_context_parallel.py
- examples/features/data_parallel/data_parallel_offline.py
- examples/offline_inference/data_parallel.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/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
- 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
- label: Distributed Tests (4xA100-4xMI300) # TBD
timeout_in_minutes: 180
@@ -1204,7 +1203,7 @@ steps:
- tests/distributed/test_torchrun_example.py
- tests/distributed/test_torchrun_example_moe.py
- examples/rl/
- tests/examples/features/data_parallel/data_parallel_offline.py
- tests/examples/offline_inference/data_parallel.py
- vllm/platforms/rocm.py
commands:
- export TORCH_NCCL_BLOCKING_WAIT=1
@@ -1214,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/features/data_parallel/data_parallel_offline.py --enforce-eager
- python3 ../examples/offline_inference/data_parallel.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
@@ -1267,7 +1266,7 @@ steps:
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- examples/features/torchrun/torchrun_dp_example_offline.py
- examples/offline_inference/torchrun_dp_example.py
- vllm/config/parallel.py
- vllm/distributed/
- vllm/v1/engine/llm_engine.py
@@ -1276,7 +1275,7 @@ steps:
- vllm/platforms/rocm.py
commands:
- export TORCH_NCCL_BLOCKING_WAIT=1
- torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
- torchrun --nproc-per-node=8 ../examples/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
#-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------#
@@ -1648,18 +1647,18 @@ steps:
- python3 basic/offline_inference/embed.py
- python3 basic/offline_inference/score.py
# Multi-modal models
- 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
- 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
# Pooling models
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# Features demo
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 offline_inference/prefix_caching.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 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
- 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
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
@@ -1952,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/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
- 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
#------------------------------------------------------- mi300 · quantization --------------------------------------------------------#
@@ -2303,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/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
- 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
- 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
@@ -2714,7 +2713,7 @@ steps:
- vllm/v1/attention/selector.py
- tests/distributed/test_context_parallel.py
- tests/v1/distributed/test_dbo.py
- examples/features/data_parallel/data_parallel_offline.py
- examples/offline_inference/data_parallel.py
- vllm/_aiter_ops.py
- vllm/platforms/rocm.py
commands:
@@ -2931,18 +2930,18 @@ steps:
- python3 basic/offline_inference/embed.py
- python3 basic/offline_inference/score.py
# Multi-modal models
- 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
- 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
# Pooling models
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# Features demo
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 offline_inference/prefix_caching.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 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
- 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
#---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------#
+1 -3
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: V1 attention (H100)
key: v1-attention-h100
timeout_in_minutes: 30
device: h100
source_file_dependencies:
@@ -15,9 +14,8 @@ steps:
- pytest -v -s v1/attention
- label: V1 attention (B200)
key: v1-attention-b200
timeout_in_minutes: 30
device: b200-k8s
device: b200
source_file_dependencies:
- vllm/config/attention.py
- vllm/model_executor/layers/attention
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Basic Correctness
key: basic-correctness
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
+1 -3
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Benchmarks CLI Test
key: benchmarks-cli-test
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
@@ -13,8 +12,7 @@ steps:
- pytest -v -s benchmarks/
- label: Attention Benchmarks Smoke Test (B200)
key: attention-benchmarks-smoke-test-b200
device: b200-k8s
device: b200
num_gpus: 2
optional: true
working_dir: "/vllm-workspace/"
+4 -17
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Sequence Parallel Correctness Tests (2 GPUs)
key: sequence-parallel-correctness-tests-2-gpus
timeout_in_minutes: 50
working_dir: "/vllm-workspace/"
num_devices: 2
@@ -18,7 +17,6 @@ steps:
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
- label: Sequence Parallel Correctness Tests (2xH100)
key: sequence-parallel-correctness-tests-2xh100
timeout_in_minutes: 50
working_dir: "/vllm-workspace/"
device: h100
@@ -29,7 +27,6 @@ steps:
- pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
- label: AsyncTP Correctness Tests (2xH100)
key: asynctp-correctness-tests-2xh100
timeout_in_minutes: 50
working_dir: "/vllm-workspace/"
device: h100
@@ -40,10 +37,9 @@ steps:
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
- label: AsyncTP Correctness Tests (B200)
key: asynctp-correctness-tests-b200
timeout_in_minutes: 50
working_dir: "/vllm-workspace/"
device: b200-k8s
device: b200
optional: true
num_devices: 2
commands:
@@ -51,7 +47,6 @@ steps:
- pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
- label: Distributed Compile Unit Tests (2xH100)
key: distributed-compile-unit-tests-2xh100
timeout_in_minutes: 20
working_dir: "/vllm-workspace/"
device: h100
@@ -65,10 +60,9 @@ steps:
- pytest -s -v tests/compile/passes/distributed
- label: Fusion and Compile Unit Tests (2xB200)
key: fusion-and-compile-unit-tests-2xb200
timeout_in_minutes: 20
working_dir: "/vllm-workspace/"
device: b200-k8s
device: b200
source_file_dependencies:
- csrc/quantization/fp4/
- vllm/model_executor/layers/quantization/
@@ -95,7 +89,6 @@ steps:
- pytest -v -s tests/compile/fullgraph/test_full_graph.py::test_fp8_kv_scale_compile
- label: Fusion E2E Quick (H100)
key: fusion-e2e-quick-h100
timeout_in_minutes: 15
working_dir: "/vllm-workspace/"
device: h100
@@ -114,7 +107,6 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and (qwen3 or deepseek)"
- label: Fusion E2E Config Sweep (H100)
key: fusion-e2e-config-sweep-h100
timeout_in_minutes: 30
working_dir: "/vllm-workspace/"
device: h100
@@ -134,10 +126,9 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "llama-3"
- label: Fusion E2E Config Sweep (B200)
key: fusion-e2e-config-sweep-b200
timeout_in_minutes: 30
working_dir: "/vllm-workspace/"
device: b200-k8s
device: b200
num_devices: 1
optional: true
commands:
@@ -148,7 +139,6 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek)) or llama-3)"
- label: Fusion E2E TP2 Quick (H100)
key: fusion-e2e-tp2-quick-h100
timeout_in_minutes: 20
working_dir: "/vllm-workspace/"
device: h100
@@ -166,7 +156,6 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
- label: Fusion E2E TP2 AR-RMS Config Sweep (H100)
key: fusion-e2e-tp2-ar-rms-config-sweep-h100
timeout_in_minutes: 40
working_dir: "/vllm-workspace/"
device: h100
@@ -186,7 +175,6 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3"
- label: Fusion E2E TP2 AsyncTP Config Sweep (H100)
key: fusion-e2e-tp2-asynctp-config-sweep-h100
timeout_in_minutes: 40
working_dir: "/vllm-workspace/"
device: h100
@@ -206,10 +194,9 @@ steps:
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "llama-3"
- label: Fusion E2E TP2 (B200)
key: fusion-e2e-tp2-b200
timeout_in_minutes: 20
working_dir: "/vllm-workspace/"
device: b200-k8s
device: b200
num_devices: 2
source_file_dependencies:
- csrc/quantization/
-2
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Platform Tests (CUDA)
key: platform-tests-cuda
timeout_in_minutes: 15
device: h200_18gb
source_file_dependencies:
@@ -14,7 +13,6 @@ steps:
- pytest -v -s cuda/test_platform_no_cuda_init.py
- label: Cudagraph
key: cudagraph
timeout_in_minutes: 20
source_file_dependencies:
- tests/v1/cudagraph
-106
View File
@@ -1,106 +0,0 @@
group: Disaggregated
depends_on:
- image-build
steps:
- label: Distributed NixlConnector PD accuracy (4 GPUs)
key: distributed-nixlconnector-pd-accuracy-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
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)
key: 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)
key: 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)
key: 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)
key: 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)
key: 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)
key: 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)
key: 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
+94 -27
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Distributed Comm Ops
key: distributed-comm-ops
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -17,7 +16,6 @@ steps:
- pytest -v -s distributed/test_shm_storage.py
- label: Distributed DP Tests (2 GPUs)
key: distributed-dp-tests-2-gpus
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -39,7 +37,6 @@ steps:
- DP_SIZE=2 pytest -v -s entrypoints/openai/test_multi_api_servers.py
- label: Distributed Compile + RPC Tests (2 GPUs)
key: distributed-compile-rpc-tests-2-gpus
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -62,7 +59,6 @@ steps:
- pytest -v -s ./compile/test_wrapper.py
- label: Distributed Torchrun + Shutdown Tests (2 GPUs)
key: distributed-torchrun-shutdown-tests-2-gpus
timeout_in_minutes: 20
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -85,7 +81,6 @@ steps:
- pytest -v -s v1/worker/test_worker_memory_snapshot.py
- label: Distributed Torchrun + Examples (4 GPUs)
key: distributed-torchrun-examples-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace"
num_devices: 4
@@ -93,8 +88,9 @@ 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/features/data_parallel/data_parallel_offline.py
- tests/examples/offline_inference/data_parallel.py
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
@@ -111,13 +107,12 @@ 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/features/data_parallel/data_parallel_offline.py --enforce-eager
- python3 examples/offline_inference/data_parallel.py --enforce-eager
# rlhf examples
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_nccl.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_ipc.py
- label: Distributed DP Tests (4 GPUs)
key: distributed-dp-tests-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -138,7 +133,6 @@ steps:
- pytest -v -s distributed/test_utils.py
- label: Distributed Compile + Comm (4 GPUs)
key: distributed-compile-comm-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -160,28 +154,24 @@ steps:
- pytest -v -s distributed/test_multiproc_executor.py::test_multiproc_executor_multi_node
- label: Distributed Tests (8 GPUs)(H100)
key: distributed-tests-8-gpus-h100
timeout_in_minutes: 10
device: h100
num_devices: 8
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- examples/features/torchrun/torchrun_dp_example_offline.py
- examples/offline_inference/torchrun_dp_example.py
- vllm/config/parallel.py
- vllm/distributed/
- vllm/v1/engine/llm_engine.py
- vllm/v1/executor/uniproc_executor.py
- vllm/v1/worker/gpu_worker.py
- tests/distributed/test_mnnvl_alltoall.py
commands:
# https://github.com/NVIDIA/nccl/issues/1838
- export NCCL_CUMEM_HOST_ENABLE=0
# test with torchrun tp=2 and dp=4 with ep
- torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
- torchrun --nproc-per-node=8 ../examples/offline_inference/torchrun_dp_example.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep
- label: Distributed Tests (4 GPUs)(A100)
key: distributed-tests-4-gpus-a100
device: a100
optional: true
num_devices: 4
@@ -196,7 +186,6 @@ steps:
- pytest -v -s -x lora/test_mixtral.py
- label: Distributed Tests (2 GPUs)(H100)
key: distributed-tests-2-gpus-h100
timeout_in_minutes: 15
device: h100
optional: true
@@ -205,14 +194,13 @@ steps:
commands:
- pytest -v -s tests/distributed/test_context_parallel.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/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
- 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
- pytest -v -s tests/v1/distributed/test_dbo.py
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 pytest -v -s tests/distributed/test_weight_transfer.py
- pytest -v -s tests/distributed/test_packed_tensor.py
- label: Distributed Tests (2 GPUs)(B200)
key: distributed-tests-2-gpus-b200
device: b200-k8s
device: b200
optional: true
working_dir: "/vllm-workspace/"
num_devices: 2
@@ -220,12 +208,8 @@ steps:
- pytest -v -s tests/distributed/test_context_parallel.py
- pytest -v -s tests/distributed/test_nccl_symm_mem_allreduce.py
- pytest -v -s tests/v1/distributed/test_dbo.py
- pytest -v -s tests/distributed/test_mnnvl_alltoall.py
- label: 2 Node Test (4 GPUs)
key: 2-node-test-4-gpus
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -238,12 +222,96 @@ steps:
- vllm/executor/
- vllm/model_executor/models/
- tests/distributed/
- tests/examples/features/data_parallel/data_parallel_offline.py
- tests/examples/offline_inference/data_parallel.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/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"
- ./.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
- label: Pipeline + Context Parallelism (4 GPUs)
key: pipeline-context-parallelism-4-gpus
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -258,7 +326,6 @@ steps:
- pytest -v -s distributed/test_pipeline_parallel.py
- label: RayExecutorV2 (4 GPUs)
key: rayexecutorv2-4-gpus
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
num_devices: 4
-16
View File
@@ -1,16 +0,0 @@
group: Docker
depends_on:
- image-build-cpu
steps:
- label: Docker Build Metadata
timeout_in_minutes: 10
device: cpu-small
source_file_dependencies:
- .buildkite/release-pipeline.yaml
- .buildkite/scripts/docker-build-metadata-args.sh
- docker/Dockerfile
- docker/Dockerfile.cpu
- docker/docker-bake.hcl
- tests/tools/test_docker_build_metadata_args.py
commands:
- pytest -v -s tools/test_docker_build_metadata_args.py
+1 -6
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: DeepSeek V2-Lite Accuracy
key: deepseek-v2-lite-accuracy
timeout_in_minutes: 60
device: h100
optional: true
@@ -13,7 +12,6 @@ steps:
- bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_ep_eplb.sh 0.25 200 8010
- label: Qwen3-30B-A3B-FP8-block Accuracy
key: qwen3-30b-a3b-fp8-block-accuracy
timeout_in_minutes: 60
device: h100
optional: true
@@ -23,9 +21,8 @@ steps:
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020
- label: Qwen3-30B-A3B-FP8-block Accuracy (B200)
key: qwen3-30b-a3b-fp8-block-accuracy-b200
timeout_in_minutes: 60
device: b200-k8s
device: b200
optional: true
num_devices: 2
working_dir: "/vllm-workspace"
@@ -33,7 +30,6 @@ steps:
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1
- label: Qwen3-30B-A3B-FP8 DP4 Async EPLB Accuracy
key: qwen3-30b-a3b-fp8-dp4-async-eplb-accuracy
timeout_in_minutes: 60
device: h100
optional: true
@@ -43,7 +39,6 @@ steps:
- bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_dp4_async_eplb.sh 0.8 200 8050
- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100)
key: deepseek-v2-lite-prefetch-offload-accuracy-h100
timeout_in_minutes: 60
device: h100
optional: true
+1 -9
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Engine
key: engine
timeout_in_minutes: 15
device: h200_18gb
source_file_dependencies:
@@ -13,12 +12,10 @@ steps:
- tests/test_config
- tests/test_logger
- tests/test_vllm_port
- tests/test_jit_monitor.py
commands:
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
- pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py
- label: Engine (1 GPU)
key: engine-1-gpu
timeout_in_minutes: 30
source_file_dependencies:
- vllm/v1/engine/
@@ -28,7 +25,6 @@ steps:
- pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py
- label: e2e Scheduling (1 GPU)
key: e2e-scheduling-1-gpu
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -38,7 +34,6 @@ steps:
- pytest -v -s v1/e2e/general/test_async_scheduling.py
- label: e2e Core (1 GPU)
key: e2e-core-1-gpu
timeout_in_minutes: 30
source_file_dependencies:
- vllm/v1/
@@ -47,7 +42,6 @@ steps:
- pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
- label: V1 e2e (2 GPUs)
key: v1-e2e-2-gpus
timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
optional: true
num_devices: 2
@@ -64,7 +58,6 @@ steps:
- image-build-amd
- label: V1 e2e (4 GPUs)
key: v1-e2e-4-gpus
timeout_in_minutes: 60 # TODO: Fix timeout after we have more confidence in the test stability
optional: true
num_devices: 4
@@ -81,7 +74,6 @@ steps:
- image-build-amd
- label: V1 e2e (4xH100)
key: v1-e2e-4xh100
timeout_in_minutes: 60
device: h100
num_devices: 4
+1 -10
View File
@@ -2,8 +2,7 @@ group: Entrypoints
depends_on:
- image-build
steps:
- label: Entrypoints Unit Tests
key: entrypoints-unit-tests
- label: Entrypoints Unit Tests
timeout_in_minutes: 10
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -14,7 +13,6 @@ steps:
- pytest -v -s entrypoints/ --ignore=entrypoints/llm --ignore=entrypoints/rpc --ignore=entrypoints/sleep --ignore=entrypoints/serve/instrumentator --ignore=entrypoints/openai --ignore=entrypoints/offline_mode --ignore=entrypoints/test_chat_utils.py --ignore=entrypoints/pooling
- label: Entrypoints Integration (LLM)
key: entrypoints-integration-llm
timeout_in_minutes: 40
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -28,7 +26,6 @@ steps:
- pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
- label: Entrypoints Integration (API Server openai - Part 1)
key: entrypoints-integration-api-server-openai-part-1
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -46,7 +43,6 @@ steps:
- label: Entrypoints Integration (API Server openai - Part 2)
key: entrypoints-integration-api-server-openai-part-2
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -64,7 +60,6 @@ steps:
- image-build-amd
- label: Entrypoints Integration (API Server openai - Part 3)
key: entrypoints-integration-api-server-openai-part-3
timeout_in_minutes: 50
device: h200_18gb
working_dir: "/vllm-workspace/tests"
@@ -77,7 +72,6 @@ steps:
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/chat_completion --ignore=entrypoints/openai/completion --ignore=entrypoints/openai/speech_to_text/ --ignore=entrypoints/openai/correctness/ --ignore=entrypoints/openai/tool_parsers/ --ignore=entrypoints/openai/responses --ignore=entrypoints/openai/test_multi_api_servers.py
- label: Entrypoints Integration (API Server 2)
key: entrypoints-integration-api-server-2
timeout_in_minutes: 130
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -92,7 +86,6 @@ steps:
- pytest -v -s tool_use
- label: Entrypoints Integration (Pooling)
key: entrypoints-integration-pooling
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -103,7 +96,6 @@ steps:
- pytest -v -s entrypoints/pooling
- label: Entrypoints Integration (Responses API)
key: entrypoints-integration-responses-api
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
@@ -113,7 +105,6 @@ steps:
- pytest -v -s entrypoints/openai/responses
- label: OpenAI API Correctness
key: openai-api-correctness
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: EPLB Algorithm
key: eplb-algorithm
timeout_in_minutes: 15
device: h200_18gb
working_dir: "/vllm-workspace/tests"
@@ -16,7 +15,6 @@ steps:
- pytest -v -s distributed/test_eplb_utils.py
- label: EPLB Execution # 17min
key: eplb-execution
timeout_in_minutes: 27
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -28,7 +26,6 @@ steps:
- pytest -v -s distributed/test_eplb_spec_decode.py
- label: Elastic EP Scaling Test
key: elastic-ep-scaling-test
timeout_in_minutes: 20
device: h100
working_dir: "/vllm-workspace/tests"
+2 -29
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: vLLM IR Tests
key: vllm-ir-tests
timeout_in_minutes: 10
device: h200_18gb
working_dir: "/vllm-workspace/"
@@ -15,7 +14,6 @@ steps:
- pytest -v -s tests/kernels/ir
- label: Kernels Core Operation Test
key: kernels-core-operation-test
timeout_in_minutes: 75
source_file_dependencies:
- csrc/
@@ -25,7 +23,6 @@ steps:
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py
- label: Kernels MiniMax Reduce RMS Test (2 GPUs)
key: kernels-minimax-reduce-rms-test-2-gpus
timeout_in_minutes: 15
num_devices: 2
device: h100
@@ -39,7 +36,6 @@ steps:
- pytest -v -s kernels/core/test_minimax_reduce_rms.py
- label: Kernels Attention Test %N
key: kernels-attention-test
timeout_in_minutes: 35
source_file_dependencies:
- csrc/attention/
@@ -53,7 +49,6 @@ steps:
parallelism: 2
- label: Kernels Quantization Test %N
key: kernels-quantization-test
timeout_in_minutes: 90
source_file_dependencies:
- csrc/quantization/
@@ -64,7 +59,6 @@ steps:
parallelism: 2
- label: Kernels MoE Test %N
key: kernels-moe-test
timeout_in_minutes: 25
source_file_dependencies:
- csrc/quantization/cutlass_w8a8/moe/
@@ -80,7 +74,6 @@ steps:
parallelism: 5
- label: Kernels Mamba Test
key: kernels-mamba-test
timeout_in_minutes: 45
source_file_dependencies:
- csrc/mamba/
@@ -89,18 +82,7 @@ steps:
commands:
- pytest -v -s kernels/mamba
- label: Kernels KDA Test
timeout_in_minutes: 20
source_file_dependencies:
- vllm/model_executor/layers/fla/ops/kda.py
- vllm/model_executor/layers/fla/ops/chunk_delta_h.py
- vllm/model_executor/layers/fla/ops/l2norm.py
- tests/kernels/test_kda.py
commands:
- pytest -v -s kernels/test_kda.py
- label: Kernels DeepGEMM Test (H100)
key: kernels-deepgemm-test-h100
timeout_in_minutes: 45
device: h100
num_devices: 1
@@ -113,19 +95,16 @@ 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)
key: kernels-b200
timeout_in_minutes: 30
working_dir: "/vllm-workspace/"
device: b200-k8s
device: b200
# optional: true
source_file_dependencies:
- csrc/quantization/fp4/
@@ -171,7 +150,6 @@ steps:
- pytest -v -s tests/models/quantization/test_nvfp4.py
- label: Kernels Helion Test
key: kernels-helion-test
timeout_in_minutes: 30
device: h100
source_file_dependencies:
@@ -183,7 +161,6 @@ steps:
- label: Kernels FP8 MoE Test (1 H100)
key: kernels-fp8-moe-test-1-h100
timeout_in_minutes: 90
device: h100
num_devices: 1
@@ -200,7 +177,6 @@ steps:
- pytest -v -s kernels/moe/test_triton_moe_ptpc_fp8.py
- label: Kernels FP8 MoE Test (2 H100s)
key: kernels-fp8-moe-test-2-h100s
timeout_in_minutes: 90
device: h100
num_devices: 2
@@ -210,9 +186,8 @@ steps:
- pytest -v -s kernels/moe/test_deepep_moe.py
- label: Kernels Fp4 MoE Test (B200)
key: kernels-fp4-moe-test-b200
timeout_in_minutes: 60
device: b200-k8s
device: b200
num_devices: 1
optional: true
commands:
@@ -223,7 +198,6 @@ steps:
- label: Kernels FusedMoE Layer Test (2 H100s)
key: kernels-fusedmoe-layer-test-2-h100s
timeout_in_minutes: 90
device: h100
num_devices: 2
@@ -240,7 +214,6 @@ steps:
- label: Kernels FusedMoE Layer Test (2 B200s)
key: kernels-fusedmoe-layer-test-2-b200s
timeout_in_minutes: 90
device: b200
num_devices: 2
+2 -20
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: LM Eval Small Models
key: lm-eval-small-models
timeout_in_minutes: 75
source_file_dependencies:
- csrc/
@@ -25,7 +24,6 @@ steps:
# - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4
- label: LM Eval Large Models (4 GPUs)(H100)
key: lm-eval-large-models-4-gpus-h100
device: h100
optional: true
num_devices: 4
@@ -38,7 +36,6 @@ steps:
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-hopper.txt --tp-size=4
- label: LM Eval Small Models (B200)
key: lm-eval-small-models-b200
timeout_in_minutes: 120
device: b200
optional: true
@@ -49,9 +46,8 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-blackwell.txt
- label: LM Eval Qwen3.5 Models (B200)
key: lm-eval-qwen3-5-models-b200
timeout_in_minutes: 120
device: b200-k8s
device: b200
optional: true
num_devices: 2
source_file_dependencies:
@@ -66,7 +62,6 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
- label: LM Eval Large Models (H200)
key: lm-eval-large-models-h200
timeout_in_minutes: 60
device: h200
optional: true
@@ -75,7 +70,6 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-h200.txt
- label: MoE Refactor Integration Test (H100 - TEMPORARY)
key: moe-refactor-integration-test-h100-temporary
device: h100
optional: true
num_devices: 2
@@ -83,15 +77,13 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-h100.txt
- label: MoE Refactor Integration Test (B200 - TEMPORARY)
key: moe-refactor-integration-test-b200-temporary
device: b200-k8s
device: b200
optional: true
num_devices: 2
commands:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt
- label: MoE Refactor Integration Test (B200 DP - TEMPORARY)
key: moe-refactor-integration-test-b200-dp-temporary
device: b200
optional: true
num_devices: 2
@@ -100,7 +92,6 @@ steps:
- label: LM Eval TurboQuant KV Cache
key: lm-eval-turboquant-kv-cache
timeout_in_minutes: 75
source_file_dependencies:
- vllm/model_executor/layers/quantization/turboquant/
@@ -111,7 +102,6 @@ steps:
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/models-turboquant.txt
- label: GPQA Eval (GPT-OSS) (H100)
key: gpqa-eval-gpt-oss-h100
timeout_in_minutes: 120
device: h100
optional: true
@@ -125,7 +115,6 @@ steps:
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt
- label: GPQA Eval (GPT-OSS) (B200)
key: gpqa-eval-gpt-oss-b200
timeout_in_minutes: 120
device: b200
optional: true
@@ -137,10 +126,3 @@ steps:
commands:
- uv pip install --system 'gpt-oss[eval]==0.0.5'
- pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
- label: MRCR Eval Small Models
timeout_in_minutes: 30
source_file_dependencies:
- tests/evals/mrcr/
commands:
- pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt
-2
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: LoRA %N
key: lora
timeout_in_minutes: 30
source_file_dependencies:
- vllm/lora
@@ -14,7 +13,6 @@ steps:
- label: LoRA TP (Distributed)
key: lora-tp-distributed
timeout_in_minutes: 30
num_devices: 4
source_file_dependencies:
+9 -22
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: V1 Spec Decode
key: v1-spec-decode
timeout_in_minutes: 30
source_file_dependencies:
- vllm/
@@ -19,7 +18,6 @@ steps:
- image-build-amd
- label: V1 Sample + Logits
key: v1-sample-logits
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -43,7 +41,6 @@ steps:
- image-build-amd
- label: V1 Core + KV + Metrics
key: v1-core-kv-metrics
timeout_in_minutes: 30
source_file_dependencies:
- vllm/
@@ -74,7 +71,6 @@ steps:
- image-build-amd
- label: V1 Others (CPU)
key: v1-others-cpu
depends_on:
- image-build-cpu
source_file_dependencies:
@@ -90,7 +86,6 @@ steps:
- pytest -v -s -m 'cpu_test' v1/metrics
- label: Regression
key: regression
timeout_in_minutes: 20
device: h200_18gb
source_file_dependencies:
@@ -102,7 +97,6 @@ steps:
working_dir: "/vllm-workspace/tests" # optional
- label: Examples
key: examples
timeout_in_minutes: 45
working_dir: "/vllm-workspace/examples"
source_file_dependencies:
@@ -119,22 +113,21 @@ steps:
- python3 basic/offline_inference/embed.py
- python3 basic/offline_inference/score.py
# for multi-modal models
- 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
- 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
# for pooling models
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# for features demo
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 offline_inference/prefix_caching.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 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 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
# 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 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
- 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
- label: Metrics, Tracing (2 GPUs)
key: metrics-tracing-2-gpus
timeout_in_minutes: 20
num_devices: 2
source_file_dependencies:
@@ -149,7 +142,6 @@ steps:
- pytest -v -s v1/tracing
- label: Python-only Installation
key: python-only-installation
depends_on: ~
timeout_in_minutes: 20
source_file_dependencies:
@@ -159,7 +151,6 @@ steps:
- bash standalone_tests/python_only_compile.sh
- label: Async Engine, Inputs, Utils, Worker
key: async-engine-inputs-utils-worker
timeout_in_minutes: 50
source_file_dependencies:
- vllm/
@@ -172,8 +163,7 @@ steps:
- pytest -v -s utils_
- label: Async Engine, Inputs, Utils, Worker, Config (CPU)
key: async-engine-inputs-utils-worker-config-cpu
depends_on:
depends_on:
- image-build-cpu
timeout_in_minutes: 30
source_file_dependencies:
@@ -206,7 +196,6 @@ steps:
- pytest -v -s config
- label: Batch Invariance (H100)
key: batch-invariance-h100
timeout_in_minutes: 30
device: h100
source_file_dependencies:
@@ -222,9 +211,8 @@ steps:
- VLLM_TEST_MODEL=Qwen/Qwen3-30B-A3B-Thinking-2507-FP8 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[FLASH_ATTN]
- label: Batch Invariance (B200)
key: batch-invariance-b200
timeout_in_minutes: 30
device: b200-k8s
device: b200
source_file_dependencies:
- vllm/v1/attention
- vllm/model_executor/layers
@@ -239,7 +227,6 @@ steps:
- pytest -v -s v1/determinism/test_nvfp4_batch_invariant.py
- label: Acceptance Length Test (Large Models) # optional
key: acceptance-length-test-large-models
timeout_in_minutes: 25
gpu: h100
optional: true
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Model Executor
key: model-executor
timeout_in_minutes: 35
source_file_dependencies:
- vllm/engine/arg_utils.py
+8 -14
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Model Runner V2 Core Tests
key: model-runner-v2-core-tests
timeout_in_minutes: 45
source_file_dependencies:
- vllm/v1/worker/gpu/
@@ -26,16 +25,14 @@ steps:
- pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0"
- label: Model Runner V2 Examples
key: model-runner-v2-examples
timeout_in_minutes: 45
working_dir: "/vllm-workspace/examples"
source_file_dependencies:
- vllm/v1/worker/gpu/
- vllm/v1/core/sched/
- vllm/v1/worker/gpu_worker.py
- examples/offline_inference/
- examples/basic/offline_inference/
- examples/generate/multimodal/
- examples/features/
- examples/pooling/embed/vision_embedding_offline.py
- examples/others/tensorize_vllm_model.py
commands:
@@ -47,22 +44,21 @@ 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 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
- 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
# for pooling models
- python3 pooling/embed/vision_embedding_offline.py --seed 0
# for features demo
- python3 features/automatic_prefix_caching/prefix_caching_offline.py
- python3 offline_inference/prefix_caching.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 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 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
# 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 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
- 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
- label: Model Runner V2 Distributed (2 GPUs)
key: model-runner-v2-distributed-2-gpus
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -83,7 +79,6 @@ steps:
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
- label: Model Runner V2 Pipeline Parallelism (4 GPUs)
key: model-runner-v2-pipeline-parallelism-4-gpus
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
num_devices: 4
@@ -99,7 +94,6 @@ steps:
- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
- label: Model Runner V2 Spec Decode
key: model-runner-v2-spec-decode
timeout_in_minutes: 30
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
+5 -11
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Basic Models Tests (Initialization)
key: basic-models-tests-initialization
timeout_in_minutes: 45
torch_nightly: true
source_file_dependencies:
@@ -17,7 +16,6 @@ steps:
torch_nightly: {}
- label: Basic Models Tests (Extra Initialization) %N
key: basic-models-tests-extra-initialization
timeout_in_minutes: 45
source_file_dependencies:
- vllm/model_executor/models/
@@ -33,7 +31,6 @@ steps:
torch_nightly: {}
- label: Basic Models Tests (Other)
key: basic-models-tests-other
timeout_in_minutes: 45
source_file_dependencies:
- vllm/
@@ -50,7 +47,6 @@ steps:
- label: Basic Models Test (Other CPU) # 5min
key: basic-models-test-other-cpu
depends_on:
- image-build-cpu
timeout_in_minutes: 10
@@ -63,7 +59,6 @@ steps:
- pytest -v -s models/test_utils.py models/test_vision.py
- label: Transformers Nightly Models
key: transformers-nightly-models
working_dir: "/vllm-workspace/"
optional: true
soft_fail: true
@@ -74,12 +69,11 @@ 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/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
# Whisper needs spawn method to avoid deadlock
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
- label: Transformers Backward Compatibility Models Test
key: transformers-backward-compatibility-models-test
working_dir: "/vllm-workspace/"
optional: true
soft_fail: true
@@ -89,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/basic/offline_inference/chat.py
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
- python3 examples/offline_inference/basic/chat.py
- python3 examples/offline_inference/vision_language.py --model-type qwen2_5_vl
# Whisper needs spawn method to avoid deadlock
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper
- VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Distributed Model Tests (2 GPUs)
key: distributed-model-tests-2-gpus
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
num_devices: 2
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Language Models Tests (Standard)
key: language-models-tests-standard
timeout_in_minutes: 25
source_file_dependencies:
- vllm/
@@ -16,7 +15,6 @@ steps:
torch_nightly: {}
- label: Language Models Tests (Extra Standard) %N
key: language-models-tests-extra-standard
timeout_in_minutes: 45
source_file_dependencies:
- vllm/model_executor/models/
@@ -33,7 +31,6 @@ steps:
torch_nightly: {}
- label: Language Models Tests (Hybrid) %N
key: language-models-tests-hybrid
timeout_in_minutes: 75
source_file_dependencies:
- vllm/
@@ -50,7 +47,6 @@ steps:
torch_nightly: {}
- label: Language Models Test (Extended Generation) # 80min
key: language-models-test-extended-generation
timeout_in_minutes: 110
optional: true
source_file_dependencies:
@@ -73,7 +69,6 @@ steps:
- pytest -v -s models/language/generation -m '(not core_model) and (not hybrid_model)'
- label: Language Models Test (PPL)
key: language-models-test-ppl
timeout_in_minutes: 110
device: h200_18gb
optional: true
@@ -84,7 +79,6 @@ steps:
- pytest -v -s models/language/generation_ppl_test
- label: Language Models Test (Extended Pooling) # 36min
key: language-models-test-extended-pooling
timeout_in_minutes: 50
optional: true
source_file_dependencies:
@@ -99,7 +93,6 @@ steps:
- image-build-amd
- label: Language Models Test (MTEB)
key: language-models-test-mteb
timeout_in_minutes: 110
device: h200_18gb
optional: true
+1 -13
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: "Multi-Modal Models (Standard) 1: qwen2"
key: multi-modal-models-standard-1-qwen2
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
@@ -20,7 +19,6 @@ steps:
- image-build-amd
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
key: multi-modal-models-standard-2-qwen3-gemma
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
@@ -30,7 +28,6 @@ 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
@@ -38,7 +35,6 @@ steps:
- image-build-amd
- label: "Multi-Modal Models (Standard) 3: llava + qwen2_vl"
key: multi-modal-models-standard-3-llava-qwen2-vl
timeout_in_minutes: 45
source_file_dependencies:
- vllm/
@@ -54,7 +50,6 @@ steps:
- image-build-amd
- label: "Multi-Modal Models (Standard) 4: other + whisper"
key: multi-modal-models-standard-4-other-whisper
timeout_in_minutes: 45
source_file_dependencies:
- vllm/
@@ -71,8 +66,7 @@ steps:
- image-build-amd
- label: Multi-Modal Processor (CPU)
key: multi-modal-processor-cpu
depends_on:
depends_on:
- image-build-cpu
timeout_in_minutes: 60
source_file_dependencies:
@@ -85,7 +79,6 @@ steps:
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
- label: Multi-Modal Processor # 44min
key: multi-modal-processor
timeout_in_minutes: 60
device: h200_18gb
source_file_dependencies:
@@ -97,7 +90,6 @@ steps:
- pytest -v -s models/multimodal/processing/test_tensor_schema.py
- label: Multi-Modal Accuracy Eval (Small Models) # 50min
key: multi-modal-accuracy-eval-small-models
timeout_in_minutes: 70
working_dir: "/vllm-workspace/.buildkite/lm-eval-harness"
source_file_dependencies:
@@ -108,7 +100,6 @@ steps:
- pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-mm-small.txt --tp-size=1
- label: Multi-Modal Models (Extended Generation 1)
key: multi-modal-models-extended-generation-1
optional: true
source_file_dependencies:
- vllm/
@@ -125,7 +116,6 @@ steps:
- image-build-amd
- label: Multi-Modal Models (Extended Generation 2)
key: multi-modal-models-extended-generation-2
optional: true
source_file_dependencies:
- vllm/
@@ -135,7 +125,6 @@ steps:
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=0) and not core_model'
- label: Multi-Modal Models (Extended Generation 3)
key: multi-modal-models-extended-generation-3
optional: true
source_file_dependencies:
- vllm/
@@ -145,7 +134,6 @@ steps:
- pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model'
- label: Multi-Modal Models (Extended Pooling)
key: multi-modal-models-extended-pooling
optional: true
device: h200_18gb
source_file_dependencies:
-1
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Plugin Tests (2 GPUs)
key: plugin-tests-2-gpus
timeout_in_minutes: 60
working_dir: "/vllm-workspace/tests"
num_devices: 2
-6
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: PyTorch Compilation Unit Tests
key: pytorch-compilation-unit-tests
timeout_in_minutes: 10
source_file_dependencies:
- vllm/
@@ -19,7 +18,6 @@ steps:
- "find compile/ -maxdepth 1 -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
- label: PyTorch Compilation Unit Tests (H100)
key: pytorch-compilation-unit-tests-h100
timeout_in_minutes: 30
device: h100
num_devices: 1
@@ -30,7 +28,6 @@ steps:
- "find compile/h100/ -name 'test_*.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
- label: PyTorch Compilation Passes Unit Tests
key: pytorch-compilation-passes-unit-tests
timeout_in_minutes: 20
source_file_dependencies:
- vllm/
@@ -39,7 +36,6 @@ steps:
- pytest -s -v compile/passes --ignore compile/passes/distributed
- label: PyTorch Fullgraph Smoke Test
key: pytorch-fullgraph-smoke-test
timeout_in_minutes: 35
source_file_dependencies:
- vllm/
@@ -52,7 +48,6 @@ steps:
- "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
- label: PyTorch Fullgraph
key: pytorch-fullgraph
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -63,7 +58,6 @@ steps:
- pytest -v -s compile/fullgraph/test_full_graph.py -k 'not test_fp8_kv_scale_compile'
- label: Pytorch Nightly Dependency Override Check # 2min
key: pytorch-nightly-dependency-override-check
# if this test fails, it means the nightly torch version is not compatible with some
# of the dependencies. Please check the error message and add the package to whitelist
# in /vllm/tools/pre_commit/generate_nightly_torch_test.py
+1 -4
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Quantization
key: quantization
timeout_in_minutes: 90
source_file_dependencies:
- csrc/
@@ -22,10 +21,9 @@ steps:
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
- label: Quantized MoE Test (B200)
key: quantized-moe-test-b200
timeout_in_minutes: 60
working_dir: "/vllm-workspace/"
device: b200-k8s
device: b200
source_file_dependencies:
- tests/quantization/test_blackwell_moe.py
- vllm/model_executor/models/deepseek_v2.py
@@ -40,7 +38,6 @@ steps:
- pytest -s -v tests/quantization/test_blackwell_moe.py
- label: Quantized Models Test
key: quantized-models-test
timeout_in_minutes: 60
source_file_dependencies:
- vllm/model_executor/layers/quantization
-1
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Ray Dependency Compatibility Check
key: ray-dependency-compatibility-check
# Informational only — does not block the pipeline.
# If this fails, it means the PR introduces a dependency that
# conflicts with Ray's dependency constraints.
+1 -4
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Samplers Test
key: samplers-test
timeout_in_minutes: 75
source_file_dependencies:
- vllm/model_executor/layers
@@ -11,9 +10,7 @@ steps:
- tests/samplers
- tests/conftest.py
commands:
# VLLM_USE_FLASHINFER_SAMPLER defaults to 1 now, so we need to pin both
# values explicitly to still cover the PyTorch-native (Triton) path.
- VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers
- pytest -v -s samplers
- VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
mirror:
amd:
+1 -22
View File
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Spec Decode Eagle
key: spec-decode-eagle
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -14,7 +13,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
- label: Spec Decode Eagle Nightly B200
key: spec-decode-eagle-nightly-b200
timeout_in_minutes: 30
device: b200
optional: true
@@ -26,7 +24,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "eagle_correctness"
- label: Spec Decode Speculators + MTP
key: spec-decode-speculators-mtp
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -38,7 +35,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- label: Spec Decode Speculators + MTP Nightly B200
key: spec-decode-speculators-mtp-nightly-b200
timeout_in_minutes: 30
device: b200
optional: true
@@ -51,7 +47,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "speculators or mtp_correctness"
- label: Spec Decode Ngram + Suffix
key: spec-decode-ngram-suffix
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -62,7 +57,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "ngram or suffix"
- label: Spec Decode Draft Model
key: spec-decode-draft-model
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
@@ -73,9 +67,8 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- label: Spec Decode Draft Model Nightly B200
key: spec-decode-draft-model-nightly-b200
timeout_in_minutes: 30
device: b200-k8s
device: b200
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
@@ -85,7 +78,6 @@ steps:
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
- label: DFlash Speculators Correctness
key: dflash-speculators-correctness
timeout_in_minutes: 30
device: h100
optional: true
@@ -97,16 +89,3 @@ steps:
commands:
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
- pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
- label: Spec Decode MTP hybrid (B200)
timeout_in_minutes: 30
device: b200
optional: true
source_file_dependencies:
- vllm/v1/spec_decode/
- vllm/v1/worker/gpu/spec_decode/
- vllm/model_executor/models/qwen3_5.py
- vllm/model_executor/models/qwen3_5_mtp.py
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
@@ -3,7 +3,6 @@ depends_on:
- image-build
steps:
- label: Weight Loading Multiple GPU # 33min
key: weight-loading-multiple-gpu
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
num_devices: 2
+9 -4
View File
@@ -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/experts/xpu_moe.py
- files=vllm/model_executor/layers/fused_moe/xpu_fused_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,7 +308,8 @@ 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/features/structured_outputs/
- files=examples/offline_inference/structured_outputs.py
- files=examples/online_serving/structured_outputs/structured_outputs.py
- files~=^tests/v1/structured_output/
- files=tests/entrypoints/llm/test_struct_output_generate.py
- files~=^vllm/v1/structured_output/
@@ -324,7 +325,7 @@ pull_request_rules:
- or:
- files~=^vllm/v1/spec_decode/
- files~=^tests/v1/spec_decode/
- files=^examples/features/speculative_decoding/
- files~=^examples/.*(spec_decode|mlpspeculator|eagle|speculation).*\.py
- 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
@@ -388,7 +389,11 @@ pull_request_rules:
- files~=^tests/entrypoints/anthropic/.*tool.*
- files~=^vllm/tool_parsers/
- files=docs/features/tool_calling.md
- files~=^examples/tool_calling/
- 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
actions:
label:
add:
+11 -2
View File
@@ -16,7 +16,11 @@ permissions:
jobs:
pre-run-check:
if: github.event_name == 'pull_request'
if: >-
github.event_name == 'pull_request' &&
(github.event.action != 'labeled' ||
github.event.label.name == 'ready' ||
github.event.label.name == 'verified')
runs-on: ubuntu-latest
steps:
- name: Check PR label and author merge count
@@ -45,7 +49,12 @@ jobs:
pre-commit:
needs: pre-run-check
if: always() && (needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
if: >-
always() &&
(github.event.action != 'labeled' ||
github.event.label.name == 'ready' ||
github.event.label.name == 'verified') &&
(needs.pre-run-check.result == 'success' || needs.pre-run-check.result == 'skipped')
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 # v6.0.1
+1 -1
View File
@@ -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.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX"
export TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 8.9 9.0+PTX"
bash tools/check_repo.sh
-1
View File
@@ -237,7 +237,6 @@ ep_kernels_workspace/
# Allow tracked library source folders under submodules (e.g., benchmarks/lib)
!vllm/benchmarks/lib/
!.buildkite/scripts/lib/
# Generated gRPC protobuf files (compiled at build time from vllm_engine.proto)
vllm/grpc/vllm_engine_pb2.py
+10 -15
View File
@@ -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" "3.14")
set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12" "3.13")
# Supported AMD GPU architectures.
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201")
set(HIP_SUPPORTED_ARCHS "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1030;gfx1100;gfx1101;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,15 +94,12 @@ 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)
# 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")
set(CUDA_SUPPORTED_ARCHS "7.5;8.0;8.6;8.7;8.9;9.0;10.0;11.0;12.0;12.1")
elseif(DEFINED CMAKE_CUDA_COMPILER_VERSION AND
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8)
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")
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")
else()
set(CUDA_SUPPORTED_ARCHS "7.0;7.5;8.0;8.6;8.7;8.9;9.0")
set(CUDA_SUPPORTED_ARCHS "7.0;7.2;7.5;8.0;8.6;8.7;8.9;9.0")
endif()
#
@@ -307,12 +304,10 @@ set(VLLM_EXT_SRC
"csrc/quantization/activation_kernels.cu"
"csrc/cuda_utils_kernels.cu"
"csrc/custom_all_reduce.cu"
"csrc/torch_bindings.cpp"
"csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
"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")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
@@ -1047,13 +1042,13 @@ endif()
set(VLLM_MOE_EXT_SRC
"csrc/moe/torch_bindings.cpp"
"csrc/moe/moe_align_sum_kernels.cu"
"csrc/moe/topk_softmax_kernels.cu"
"csrc/moe/topk_softplus_sqrt_kernels.cu")
"csrc/moe/topk_softmax_kernels.cu")
if(VLLM_GPU_LANG STREQUAL "CUDA")
list(APPEND VLLM_MOE_EXT_SRC
"csrc/moe/moe_wna16.cu"
"csrc/moe/grouped_topk_kernels.cu")
"csrc/moe/grouped_topk_kernels.cu"
"csrc/moe/router_gemm.cu")
endif()
if(VLLM_GPU_LANG STREQUAL "CUDA")
+2 -2
View File
@@ -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 Hugging Face, including:
vLLM seamlessly supports 200+ model architectures on HuggingFace, including:
- Decoder-only LLMs (e.g., Llama, Qwen, Gemma)
- Mixture-of-Expert LLMs (e.g., Mixtral, DeepSeek-V3, Qwen-MoE, GPT-OSS)
@@ -404,7 +404,6 @@ 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.experts.cutlass_moe import CutlassExpertsFp8
from vllm.model_executor.layers.fused_moe.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.experts.cutlass_moe import (
from vllm.model_executor.layers.fused_moe.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.experts.cutlass_moe import CutlassExpertsFp8
from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
from vllm.model_executor.layers.fused_moe.fused_moe import (
fused_experts,
fused_topk,
@@ -1,324 +0,0 @@
# 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,
)
@@ -217,7 +217,6 @@ async def send_request(
min_tokens: int | None = None,
max_tokens: int | None = None,
timeout_sec: int = 120,
conversation_id: str | None = None,
) -> ServerResponse:
payload = {
"model": model,
@@ -226,9 +225,6 @@ async def send_request(
"temperature": 0.0,
}
if conversation_id is not None:
payload["conversation_id"] = conversation_id
if stream:
payload["stream"] = True
payload["stream_options"] = {"include_usage": False}
@@ -423,7 +419,6 @@ async def send_turn(
min_tokens,
max_tokens,
req_args.timeout_sec,
conversation_id=conv_id,
)
if response.valid is False:
@@ -1473,12 +1468,6 @@ async def main() -> None:
"(for example: --warmup-percentages=0%%,50%%)",
)
parser.add_argument(
"--trust-remote-code",
action="store_true",
help="Trust remote code when loading the tokenizer.",
)
args = parser.parse_args()
logger.info(args)
@@ -1521,9 +1510,7 @@ async def main() -> None:
np.random.seed(args.seed)
logger.info("Loading tokenizer")
tokenizer = AutoTokenizer.from_pretrained(
args.model, trust_remote_code=args.trust_remote_code
)
tokenizer = AutoTokenizer.from_pretrained(args.model)
await get_server_info(args.url)
+1 -6
View File
@@ -20,7 +20,7 @@ else()
FetchContent_Declare(
deepgemm
GIT_REPOSITORY https://github.com/deepseek-ai/DeepGEMM.git
GIT_TAG 891d57b4db1071624b5c8fa0d1e51cb317fa709f
GIT_TAG 477618cd51baffca09c4b0b87e97c03fe827ef03
GIT_SUBMODULES "third-party/cutlass" "third-party/fmt"
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
@@ -120,11 +120,6 @@ 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")
+1 -1
View File
@@ -19,7 +19,7 @@ else()
FetchContent_Declare(
flashmla
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
GIT_TAG a6ec2ba7bd0a7dff98b3f4d3e6b52b159c48d78b
GIT_TAG 692917b1cda61b93ac9ee2d846ec54e75afe87b1
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
+25 -82
View File
@@ -11,74 +11,29 @@
namespace vllm {
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
bool act_first, bool HAS_CLAMP>
bool act_first>
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
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);
}
const scalar_t& y) {
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
}
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
bool act_first, bool HAS_CLAMP>
bool act_first>
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
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));
}
const packed_t& y) {
return act_first ? packed_mul(PACKED_ACT_FN(x), y)
: packed_mul(x, PACKED_ACT_FN(y));
}
// 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 HAS_CLAMP, bool use_256b = false>
bool use_vec, 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 float limit) {
const int d) {
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;
@@ -103,9 +58,8 @@ __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, HAS_CLAMP>(
x.elts[j], y.elts[j], limit);
x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
x.elts[j], y.elts[j]);
}
if constexpr (use_256b) {
st256(x, &out_vec[i]);
@@ -118,8 +72,7 @@ __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, HAS_CLAMP>(x, y, limit);
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
}
}
}
@@ -198,11 +151,8 @@ 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. 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) \
// first.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
auto dtype = input.scalar_type(); \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
@@ -227,8 +177,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, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
} else { \
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
@@ -236,8 +186,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, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
} \
} else { \
@@ -247,8 +197,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, HAS_CLAMP><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
ACT_FIRST, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
}
@@ -256,14 +206,7 @@ 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, 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);
true);
}
void mul_and_silu(torch::Tensor& out, // [..., d]
@@ -272,21 +215,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, 0.0f);
false);
}
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, false, 0.0f);
true);
}
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, false, 0.0f);
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
vllm::packed_gelu_tanh_kernel, true);
}
namespace vllm {
+8 -18
View File
@@ -599,11 +599,6 @@ __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;
@@ -616,18 +611,16 @@ __global__ void cp_gather_indexer_k_quant_cache_kernel(
}
}
__syncthreads();
#ifndef USE_ROCM
__syncwarp();
#endif
// 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) {
if (head_idx >= head_dim || token_idx >= num_tokens) {
return;
}
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 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 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;
@@ -1497,9 +1490,6 @@ 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;
@@ -1511,7 +1501,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_HALF_TYPES(ql_nope.scalar_type(), "concat_mla_q", [&] {
VLLM_DISPATCH_FLOATING_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),
+2 -48
View File
@@ -61,23 +61,8 @@
#endif
#ifdef __aarch64__
// Implementation of neon_expf copied from Arm Optimized Routines (expf
// AdvSIMD)
// Implementation 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); \
@@ -121,38 +106,7 @@
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__
+25 -82
View File
@@ -1,16 +1,5 @@
#include "cpu_attn_dispatch_generated.h"
// Maps kv_cache_dtype string to Fp8KVCacheDataType enum.
// "auto" -> kAuto(0); "fp8"/"fp8_e4m3" -> kFp8E4M3; "fp8_e5m2" -> kFp8E5M2.
static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
const std::string& kv_cache_dtype) {
if (kv_cache_dtype == "fp8_e5m2")
return cpu_attention::Fp8KVCacheDataType::kFp8E5M2;
if (kv_cache_dtype == "fp8_e4m3" || kv_cache_dtype == "fp8")
return cpu_attention::Fp8KVCacheDataType::kFp8E4M3;
return cpu_attention::Fp8KVCacheDataType::kAuto;
}
torch::Tensor get_scheduler_metadata(
const int64_t num_req, const int64_t num_heads_q,
const int64_t num_heads_kv, const int64_t head_dim,
@@ -29,8 +18,6 @@ torch::Tensor get_scheduler_metadata(
isa = cpu_attention::ISA::NEON;
} else if (isa_hint == "vxe") {
isa = cpu_attention::ISA::VXE;
} else if (isa_hint == "vsx") {
isa = cpu_attention::ISA::VSX;
} else {
TORCH_CHECK(false, "Unsupported CPU attention ISA hint: " + isa_hint);
}
@@ -62,7 +49,7 @@ torch::Tensor get_scheduler_metadata(
input.enable_kv_split = enable_kv_split;
VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
CPU_ATTN_DISPATCH(head_dim, isa, [&]() {
input.elem_size = sizeof(scalar_t);
input.q_buffer_elem_size = sizeof(attn_impl::q_buffer_t);
input.logits_buffer_elem_size = sizeof(attn_impl::logits_buffer_t);
@@ -85,9 +72,7 @@ void cpu_attn_reshape_and_cache(
key_cache, // [num_blocks, num_kv_heads, block_size, head_size]
torch::Tensor&
value_cache, // [num_blocks, num_kv_heads, block_size, head_size]
const torch::Tensor& slot_mapping, const std::string& isa,
const double k_scale = 1.0, const double v_scale = 1.0,
const std::string& kv_cache_dtype = "auto") {
const torch::Tensor& slot_mapping, const std::string& isa) {
TORCH_CHECK_EQ(key.dim(), 3);
TORCH_CHECK_EQ(value.dim(), 3);
TORCH_CHECK_EQ(key_cache.dim(), 4);
@@ -95,30 +80,18 @@ void cpu_attn_reshape_and_cache(
TORCH_CHECK_EQ(key.stride(2), 1);
TORCH_CHECK_EQ(value.stride(2), 1);
const int64_t kv_cache_idx =
static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
const bool is_fp8 = (kv_cache_idx != 0);
if (is_fp8) {
TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
"key_cache must be uint8 for FP8 path");
TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
"value_cache must be uint8 for FP8 path");
TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
}
const float k_inv = is_fp8 ? 1.0f / static_cast<float>(k_scale) : 0.0f;
const float v_inv = is_fp8 ? 1.0f / static_cast<float>(v_scale) : 0.0f;
const int64_t token_num = key.size(0);
const int64_t head_num = key.size(1);
const int64_t head_dim = key.size(2);
const int64_t key_token_num_stride = key.stride(0);
const int64_t value_token_num_stride = value.stride(0);
const int64_t head_num = value.size(1);
const int64_t key_head_num_stride = key.stride(1);
const int64_t value_head_num_stride = value.stride(1);
const int64_t num_blocks = key_cache.size(0);
const int64_t num_blocks_stride = key_cache.stride(0);
const int64_t cache_head_num_stride = key_cache.stride(1);
const int64_t block_size = key_cache.size(2);
const int64_t block_size_stride = key_cache.stride(2);
const int64_t head_dim = key.size(-1);
cpu_attention::ISA isa_tag = [&]() {
if (isa == "amx") {
@@ -131,31 +104,21 @@ void cpu_attn_reshape_and_cache(
return cpu_attention::ISA::NEON;
} else if (isa == "vxe") {
return cpu_attention::ISA::VXE;
} else if (isa == "vsx") {
return cpu_attention::ISA::VSX;
} else {
TORCH_CHECK(false, "Invalid ISA type: " + isa);
}
}();
if (is_fp8) {
TORCH_CHECK(isa_tag == cpu_attention::ISA::AMX ||
isa_tag == cpu_attention::ISA::VEC,
"FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
}
VLLM_DISPATCH_FLOATING_TYPES(
key.scalar_type(), "cpu_attn_reshape_and_cache", [&]() {
CPU_ATTN_DISPATCH(head_dim, isa_tag, kv_cache_idx, [&]() {
using kv_t = typename attn_impl::kv_cache_t;
CPU_ATTN_DISPATCH(head_dim, isa_tag, [&]() {
attn_impl::reshape_and_cache(
key.data_ptr<scalar_t>(), value.data_ptr<scalar_t>(),
reinterpret_cast<kv_t*>(key_cache.data_ptr()),
reinterpret_cast<kv_t*>(value_cache.data_ptr()),
slot_mapping.data_ptr<int64_t>(), token_num, key.stride(0),
value.stride(0), head_num, key.stride(1), value.stride(1),
num_blocks, num_blocks_stride, cache_head_num_stride, block_size,
block_size_stride, k_inv, v_inv);
key_cache.data_ptr<scalar_t>(), value_cache.data_ptr<scalar_t>(),
slot_mapping.data_ptr<int64_t>(), token_num, key_token_num_stride,
value_token_num_stride, head_num, key_head_num_stride,
value_head_num_stride, num_blocks, num_blocks_stride,
cache_head_num_stride, block_size, block_size_stride);
});
});
}
@@ -174,26 +137,13 @@ void cpu_attention_with_kv_cache(
const int64_t sliding_window_left, const int64_t sliding_window_right,
const torch::Tensor& block_table, // [num_tokens, max_block_num]
const double softcap, const torch::Tensor& scheduler_metadata,
const std::optional<torch::Tensor>& s_aux, // [num_heads]
const double k_scale = 1.0, const double v_scale = 1.0,
const std::string& kv_cache_dtype = "auto") {
const std::optional<torch::Tensor>& s_aux // [num_heads]
) {
TORCH_CHECK_EQ(query.dim(), 3);
TORCH_CHECK_EQ(query.stride(2), 1);
TORCH_CHECK_EQ(key_cache.dim(), 4);
TORCH_CHECK_EQ(value_cache.dim(), 4);
const int64_t kv_cache_idx =
static_cast<int64_t>(parse_fp8_kv_dtype(kv_cache_dtype));
const bool is_fp8 = (kv_cache_idx != 0);
if (is_fp8) {
TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte,
"key_cache must be uint8 for FP8 path");
TORCH_CHECK(value_cache.scalar_type() == at::ScalarType::Byte,
"value_cache must be uint8 for FP8 path");
TORCH_CHECK(k_scale > 0, "k_scale must be positive for FP8 path");
TORCH_CHECK(v_scale > 0, "v_scale must be positive for FP8 path");
}
cpu_attention::AttentionInput input;
input.metadata = reinterpret_cast<cpu_attention::AttentionMetadata*>(
scheduler_metadata.data_ptr());
@@ -215,32 +165,25 @@ void cpu_attention_with_kv_cache(
input.block_table = block_table.data_ptr<int32_t>();
input.alibi_slopes =
alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
// For now sink must be bf16
input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
input.scale = scale;
input.causal = causal;
input.sliding_window_left = sliding_window_left;
input.sliding_window_right = sliding_window_right;
if (input.causal) {
// to make boundary calculation easier
input.sliding_window_right = 0;
}
input.softcap = static_cast<float>(softcap);
if (is_fp8) {
input.k_scale_fp8 = static_cast<float>(k_scale);
input.v_scale_fp8 = static_cast<float>(v_scale);
TORCH_CHECK(input.metadata->isa == cpu_attention::ISA::AMX ||
input.metadata->isa == cpu_attention::ISA::VEC,
"FP8 KV cache is only supported on x86 (AMX/VEC) ISA");
}
float softcap_fp32 = softcap;
input.softcap = softcap_fp32;
VLLM_DISPATCH_FLOATING_TYPES(
query.scalar_type(), "cpu_attention_with_kv_cache", [&]() {
CPU_ATTN_DISPATCH(
query.size(2), input.metadata->isa, kv_cache_idx, [&]() {
TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment,
0);
cpu_attention::AttentionMainLoop<attn_impl> mainloop;
mainloop(&input);
});
CPU_ATTN_DISPATCH(query.size(2), input.metadata->isa, [&]() {
TORCH_CHECK_EQ(input.block_size % attn_impl::BlockSizeAlignment, 0);
cpu_attention::AttentionMainLoop<attn_impl> mainloop;
mainloop(&input);
});
});
}
+46 -171
View File
@@ -1,7 +1,6 @@
#ifndef CPU_ATTN_AMX_HPP
#define CPU_ATTN_AMX_HPP
#include "cpu_attn_fp8.hpp"
#include "cpu_attn_impl.hpp"
namespace cpu_attention {
@@ -22,10 +21,9 @@ typedef struct __tile_config {
// 2-2-4 pattern, for 16 < m <= 32
// TILE 0, 1: load A matrix, row num should be 16, m - 16
// TILE 2, 3: load B matrix, row num should be 16
// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16,
// m - 16, m - 16
// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
template <typename q_buffer_t, typename kv_cache_t>
// TILE 4, 5, 6, 7: store results C matrix, row num should be 16, 16, m - 16, m
// - 16
template <typename kv_cache_t>
class TileGemm224 {
public:
template <AttentionGemmPhase phase, int32_t k_size>
@@ -44,56 +42,13 @@ class TileGemm224 {
}
};
// Dequantize one FP8 tile (AMX_TILE_ROW_NUM rows x 32 cols) to BF16.
template <typename kv_cache_t>
FORCE_INLINE void deq_tile_amx(const uint8_t* src, c10::BFloat16* dst) {
for (int r = 0; r < AMX_TILE_ROW_NUM; ++r) {
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e4m3_tag{})
.save(dst + r * 32);
} else {
vec_op::BF16Vec32(src + r * 32, vec_op::fp8_bf16_e5m2_tag{})
.save(dst + r * 32);
}
}
}
// For FP8: dequant src into scratch and return scratch.
// For BF16: return src directly (scratch is unused; the compiler elides it).
template <typename kv_cache_t>
FORCE_INLINE const c10::BFloat16* prepare_b_tile(const kv_cache_t* src,
c10::BFloat16* scratch) {
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
deq_tile_amx<kv_cache_t>(reinterpret_cast<const uint8_t*>(src), scratch);
return scratch;
} else {
return reinterpret_cast<const c10::BFloat16*>(src);
}
}
// Handles both BF16 and FP8 KV cache (2-2-4 pattern).
template <typename kv_cache_t>
class TileGemm224<c10::BFloat16, kv_cache_t> {
static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
"kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
static constexpr bool fp8_kv =
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
// BF16 path: scratch_elems=1 so the scratch array is eliminated by the
// compiler.
static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
template <>
class TileGemm224<c10::BFloat16> {
public:
template <AttentionGemmPhase phase, int32_t k_size>
FORCE_INLINE static void gemm(const int32_t m_size,
c10::BFloat16* __restrict__ a_tile,
kv_cache_t* __restrict__ b_tile,
c10::BFloat16* __restrict__ b_tile,
float* __restrict__ c_tile, const int64_t lda,
const int64_t ldb, const int64_t ldc,
const int32_t block_size,
@@ -101,7 +56,6 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
const bool accum_c) {
const int32_t k_times =
dynamic_k_size / (AMX_TILE_ROW_NUM * 4 / sizeof(c10::BFloat16));
c10::BFloat16* __restrict__ a_tile_0 = a_tile;
c10::BFloat16* __restrict__ a_tile_1 = a_tile + lda * AMX_TILE_ROW_NUM;
const int64_t a_tile_stride = [&]() {
@@ -116,8 +70,8 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
}
}();
kv_cache_t* __restrict__ b_tile_2 = b_tile;
kv_cache_t* __restrict__ b_tile_3 = [&]() {
c10::BFloat16* __restrict__ b_tile_2 = b_tile;
c10::BFloat16* __restrict__ b_tile_3 = [&]() {
if constexpr (phase == AttentionGemmPhase::QK) {
// k_cache is prepacked
return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
@@ -152,16 +106,11 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
_tile_zero(7);
}
alignas(64) c10::BFloat16 scratch_2[scratch_elems];
alignas(64) c10::BFloat16 scratch_3[scratch_elems];
for (int32_t k = 0; k < k_times; ++k) {
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
_tile_loadd(0, a_tile_0, a_tile_stride);
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_tile_stride);
_tile_stream_loadd(2, b_tile_2, b_tile_stride);
_tile_dpbf16ps(4, 0, 2);
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_tile_stride);
_tile_stream_loadd(3, b_tile_3, b_tile_stride);
_tile_dpbf16ps(5, 0, 3);
_tile_loadd(1, a_tile_1, a_tile_stride);
_tile_dpbf16ps(6, 1, 2);
@@ -205,13 +154,13 @@ class TileGemm224<c10::BFloat16, kv_cache_t> {
};
// 1-2-2 pattern, for 0 < m <= 16
// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should
// be m, m
// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row num
// should be 16
// TILE 6, 7: store results C matrix, row num should be m
// q_buffer_t: A (Q/P) tile type; kv_cache_t: B (K/V cache) tile type.
template <typename q_buffer_t, typename kv_cache_t>
// TILE 0, (1): load A matrix, use extra 1 tile for prefetch, row num should be
// m, m
// TILE 2, 3, (4, 5): load B matrix, use extra 2 tiles for prefetch, row
// num should be 16
// TILE 6, 7, (6, 7): store results C matrix, row num should be
// m
template <typename kv_cache_t>
class TileGemm122 {
public:
template <AttentionGemmPhase phase, int32_t k_size>
@@ -230,26 +179,13 @@ class TileGemm122 {
}
};
// Handles both BF16 and FP8 KV cache (1-2-2 pattern).
template <typename kv_cache_t>
class TileGemm122<c10::BFloat16, kv_cache_t> {
static_assert(std::is_same_v<kv_cache_t, c10::BFloat16> ||
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>,
"kv_cache_t must be BFloat16, Float8_e4m3fn, or Float8_e5m2");
static constexpr bool fp8_kv =
std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
static constexpr int64_t tile_elems = AMX_TILE_BYTES / sizeof(c10::BFloat16);
static constexpr int64_t scratch_elems = fp8_kv ? tile_elems : 1;
template <>
class TileGemm122<c10::BFloat16> {
public:
template <AttentionGemmPhase phase, int32_t k_size>
FORCE_INLINE static void gemm(const int32_t m_size,
c10::BFloat16* __restrict__ a_tile,
kv_cache_t* __restrict__ b_tile,
c10::BFloat16* __restrict__ b_tile,
float* __restrict__ c_tile, const int64_t lda,
const int64_t ldb, const int64_t ldc,
const int32_t block_size,
@@ -279,19 +215,21 @@ class TileGemm122<c10::BFloat16, kv_cache_t> {
}
}();
kv_cache_t* __restrict__ b_tile_2 = b_tile;
kv_cache_t* __restrict__ b_tile_3 = [&]() {
c10::BFloat16* __restrict__ b_tile_2 = b_tile;
c10::BFloat16* __restrict__ b_tile_3 = [&]() {
if constexpr (phase == AttentionGemmPhase::QK) {
// k_cache is prepacked
return b_tile + (k_size * AMX_TILE_ROW_BYTES / 4);
} else if constexpr (phase == AttentionGemmPhase::PV) {
// v_cache is prepacked
return b_tile + (block_size * AMX_TILE_ROW_BYTES / 4);
} else {
TORCH_CHECK(false, "Unreachable");
}
}();
kv_cache_t* __restrict__ b_tile_4 =
c10::BFloat16* __restrict__ b_tile_4 =
b_tile_2 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
kv_cache_t* __restrict__ b_tile_5 =
c10::BFloat16* __restrict__ b_tile_5 =
b_tile_3 + AMX_TILE_BYTES / sizeof(c10::BFloat16);
int64_t b_stride = AMX_TILE_ROW_BYTES;
@@ -312,25 +250,16 @@ class TileGemm122<c10::BFloat16, kv_cache_t> {
_tile_zero(7);
}
alignas(64) c10::BFloat16 scratch_2[scratch_elems];
alignas(64) c10::BFloat16 scratch_3[scratch_elems];
alignas(64) c10::BFloat16 scratch_4[scratch_elems];
alignas(64) c10::BFloat16 scratch_5[scratch_elems];
for (int32_t k = 0; k < k_group_times; ++k) {
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
const c10::BFloat16* load_4 = prepare_b_tile(b_tile_4, scratch_4);
const c10::BFloat16* load_5 = prepare_b_tile(b_tile_5, scratch_5);
_tile_loadd(0, a_tile_0, a_tile_stride);
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
_tile_stream_loadd(2, b_tile_2, b_stride);
_tile_dpbf16ps(6, 0, 2);
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
_tile_stream_loadd(3, b_tile_3, b_stride);
_tile_dpbf16ps(7, 0, 3);
_tile_loadd(1, a_tile_1, a_tile_stride);
_tile_stream_loadd(4, const_cast<c10::BFloat16*>(load_4), b_stride);
_tile_stream_loadd(4, b_tile_4, b_stride);
_tile_dpbf16ps(6, 1, 4);
_tile_stream_loadd(5, const_cast<c10::BFloat16*>(load_5), b_stride);
_tile_stream_loadd(5, b_tile_5, b_stride);
_tile_dpbf16ps(7, 1, 5);
// update ptrs
@@ -350,13 +279,10 @@ class TileGemm122<c10::BFloat16, kv_cache_t> {
}
if (has_tail) {
const c10::BFloat16* load_2 = prepare_b_tile(b_tile_2, scratch_2);
const c10::BFloat16* load_3 = prepare_b_tile(b_tile_3, scratch_3);
_tile_loadd(0, a_tile_0, a_tile_stride);
_tile_stream_loadd(2, const_cast<c10::BFloat16*>(load_2), b_stride);
_tile_stream_loadd(2, b_tile_2, b_stride);
_tile_dpbf16ps(6, 0, 2);
_tile_stream_loadd(3, const_cast<c10::BFloat16*>(load_3), b_stride);
_tile_stream_loadd(3, b_tile_3, b_stride);
_tile_dpbf16ps(7, 0, 3);
}
@@ -376,25 +302,21 @@ class TileGemm122<c10::BFloat16, kv_cache_t> {
_tile_loadconfig(&config);
}
};
} // namespace
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
static constexpr bool fp8_kv =
std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
template <typename scalar_t, int64_t head_dim>
class AttentionImpl<ISA::AMX, scalar_t, head_dim> {
public:
using query_t = scalar_t;
using q_buffer_t = scalar_t;
using kv_cache_t = kv_cache_scalar_t;
using kv_cache_t = scalar_t;
using logits_buffer_t = float;
using partial_output_buffer_t = float;
using prob_buffer_t = scalar_t;
constexpr static int64_t BlockSizeAlignment =
32; // AMX_TILE_ROW_NUM = 16 tokens/tile; 32 = 2 tiles
AMX_TILE_ROW_BYTES /
sizeof(kv_cache_t); // KV token num unit of QK and PV phases
constexpr static int64_t HeadDimAlignment =
2 * (AMX_TILE_ROW_BYTES / 4); // headdim num unit of PV phase
constexpr static int64_t MaxQHeadNumPerIteration = 32;
@@ -402,9 +324,6 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
constexpr static ISA ISAType = ISA::AMX;
constexpr static bool scale_on_logits = true;
float k_scale = 1.0f;
float v_scale = 1.0f;
public:
AttentionImpl() : current_q_head_num_(0) {
// Use all columns in AMX tiles
@@ -413,50 +332,21 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
~AttentionImpl() { _tile_release(); }
void init_from_input(const AttentionInput* input) {
if constexpr (fp8_kv) {
k_scale = input->k_scale_fp8;
v_scale = input->v_scale_fp8;
}
}
float get_output_v_scale() const noexcept {
if constexpr (fp8_kv) {
// AMX dequant places FP8 payload into a BF16 field (exponent bias 127).
// Correction = 2^(127 - FP8_bias): E4M3 bias=7 → 2^120, E5M2 bias=15 →
// 2^112.
constexpr float bias =
std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
return v_scale * bias;
}
return 1.0f;
}
template <template <typename tile_gemm_t> typename attention>
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
if constexpr (fp8_kv) {
// Same bias correction as get_output_v_scale: AMX FP8→BF16 dequant
// shifts the exponent bias from FP8 to BF16 (127), so we multiply by
// 2^(127-FP8_bias) to recover the true value. E4M3: 2^120, E5M2: 2^112.
const float bias =
std::is_same_v<kv_cache_t, c10::Float8_e5m2> ? 0x1p112f : 0x1p120f;
scale *= k_scale * bias;
}
if (q_head_num > AMX_TILE_ROW_NUM) {
if (q_head_num != current_q_head_num_) {
current_q_head_num_ = q_head_num;
TileGemm224<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
amx_tile_config_);
TileGemm224<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
}
attention<TileGemm224<q_buffer_t, kv_cache_t>> attention_iteration;
attention<TileGemm224<kv_cache_t>> attention_iteration;
attention_iteration(CPU_ATTENTION_PARAMS);
} else {
if (q_head_num != current_q_head_num_) {
current_q_head_num_ = q_head_num;
TileGemm122<q_buffer_t, kv_cache_t>::init_tile_config(q_head_num,
amx_tile_config_);
TileGemm122<kv_cache_t>::init_tile_config(q_head_num, amx_tile_config_);
}
attention<TileGemm122<q_buffer_t, kv_cache_t>> attention_iteration;
attention<TileGemm122<kv_cache_t>> attention_iteration;
attention_iteration(CPU_ATTENTION_PARAMS);
}
}
@@ -521,26 +411,13 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
// reshape KV to AMX friendly layout
static void reshape_and_cache(
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
const int64_t head_num, const int64_t key_head_num_stride,
const int64_t value_head_num_stride, const int64_t num_blocks,
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
const int64_t block_size, const int64_t block_size_stride,
const float k_inv = 0.0f, const float v_inv = 0.0f) {
if constexpr (fp8_kv) {
constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
reshape_and_cache_fp8_amx_impl<scalar_t, qfn>(
key, value, reinterpret_cast<uint8_t*>(key_cache),
reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
head_num, head_dim, block_size, key_token_num_stride,
key_head_num_stride, value_token_num_stride, value_head_num_stride,
num_blocks_stride, cache_head_num_stride, num_blocks_stride,
cache_head_num_stride, k_inv, v_inv);
return;
}
const int64_t block_size, const int64_t block_size_stride) {
// For AMX 2D tiles, size of each line is 64 bytes
constexpr int64_t amx_tile_row_size = AMX_TILE_ROW_BYTES;
// For AMX B matrix, N always is 16
@@ -549,9 +426,6 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
// For now suppose block_size is divisible by amx_tile_column_num
TORCH_CHECK_EQ(block_size % amx_b_tile_k_size, 0);
scalar_t* __restrict__ kc = reinterpret_cast<scalar_t*>(key_cache);
scalar_t* __restrict__ vc = reinterpret_cast<scalar_t*>(value_cache);
#pragma omp parallel for collapse(2)
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -579,7 +453,8 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
constexpr int64_t quadword_num_per_group =
token_num_per_group * quadword_num;
int32_t* key_cache_start_ptr =
reinterpret_cast<int32_t*>(kc + block_idx * num_blocks_stride +
reinterpret_cast<int32_t*>(key_cache +
block_idx * num_blocks_stride +
head_idx * cache_head_num_stride) +
group_idx * quadword_num_per_group + group_offset;
@@ -608,7 +483,7 @@ class AttentionImpl<ISA::AMX, scalar_t, head_dim, kv_cache_scalar_t> {
token_idx * value_token_num_stride +
head_idx * value_head_num_stride;
scalar_t* value_cache_start_ptr =
vc + block_idx * num_blocks_stride +
value_cache + block_idx * num_blocks_stride +
head_idx * cache_head_num_stride +
sub_group_idx * token_num_per_sub_group * amx_b_tile_n_size +
sub_group_offset;
-214
View File
@@ -1,214 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#pragma once
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <limits>
#include <type_traits>
#include "cpu/utils.hpp"
typedef uint32_t __attribute__((__may_alias__)) u32_alias_t;
typedef uint16_t __attribute__((__may_alias__)) u16_alias_t;
typedef float __attribute__((__may_alias__)) f32_alias_t;
// Reference scalar dequant — used to verify vectorized AMX dequant.
inline float fp8e4m3_to_float_scalar(uint8_t b, float scale) noexcept {
// NaN encoding in E4M3
if ((b & 0x7F) == 0x7F) return std::numeric_limits<float>::quiet_NaN();
uint32_t b_u32 = static_cast<uint32_t>(b);
uint32_t sign = (b_u32 & 0x80) << 24;
uint32_t payload = (b_u32 & 0x7F) << 20;
uint32_t bits = sign | payload;
float b_f32_unscaled = *reinterpret_cast<const f32_alias_t*>(&bits);
float b_f32_scaled = b_f32_unscaled * scale * 0x1p120f;
return b_f32_scaled;
}
inline uint8_t float_to_fp8e4m3_scalar(float v, float inv_scale) noexcept {
v *= inv_scale;
constexpr float fp8_max = 448.0f;
v = std::max(-fp8_max, std::min(fp8_max, v));
if (v == 0.0f) return 0;
// Inverse mapping of fp8e4m3_to_float_scalar: shift the effective exponent
// bias from fp32 (127) back to fp8 e4m3 (7), then pack sign|payload.
float v_f32_unscaled = v * 0x1p-120f;
uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v_f32_unscaled);
uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
uint8_t payload = static_cast<uint8_t>((bits >> 20) & 0x7F);
if (payload == 0) return sign;
payload = std::min<uint8_t>(payload, 0x7E); // keep 0x7F as NaN encoding
return static_cast<uint8_t>(sign | payload);
}
// ---------------------------------------------------------------------------
// AMX reshape impl — parameterised on the quantisation function.
// Writes key/value into uint8 FP8 KV cache using the AMX tile-friendly layout.
// K: halfword-packed (2 FP8 per uint16, token_num_per_group=16).
// V: sub-group packing (token_num_per_sub_group=2, head_elems_per_group=16).
// block_size must be divisible by 32.
// ---------------------------------------------------------------------------
template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
inline void reshape_and_cache_fp8_amx_impl(
const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
float v_inv) {
constexpr int64_t token_num_per_group = 16; // AMX_TILE_ROW_NUM
const int64_t halfword_num = head_dim / 2; // 2 FP8 per uint16
const int64_t halfword_num_per_group = token_num_per_group * halfword_num;
constexpr int64_t head_elems_per_group = 16;
constexpr int64_t token_num_per_sub_group = 2; // = 4 / sizeof(BF16)
const int64_t group_num = head_dim / head_elems_per_group;
const int64_t group_size = block_size * head_elems_per_group;
#pragma omp parallel for collapse(2) schedule(static)
for (int64_t tok = 0; tok < token_num; ++tok) {
for (int64_t h = 0; h < head_num; ++h) {
const int64_t slot = slot_ptr[tok];
if (slot < 0) continue;
const int64_t block_idx = slot / block_size;
const int64_t block_offset = slot % block_size;
// Key: halfword-packed, 2 FP8 per uint16
{
const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
const int64_t group_idx = block_offset / token_num_per_group;
const int64_t group_offset = block_offset % token_num_per_group;
uint16_t* kdst =
reinterpret_cast<uint16_t*>(key_cache_ptr + block_idx * kc_stride0 +
h * kc_stride1) +
group_idx * halfword_num_per_group + group_offset;
for (int64_t j = 0; j < halfword_num; ++j) {
uint8_t fp8_0 = quant_fn(static_cast<float>(ksrc[j * 2]), k_inv);
uint8_t fp8_1 = quant_fn(static_cast<float>(ksrc[j * 2 + 1]), k_inv);
uint8_t bytes[2] = {fp8_0, fp8_1};
uint16_t hw = *reinterpret_cast<const u16_alias_t*>(bytes);
kdst[j * token_num_per_group] = hw;
}
}
// Value: sub-group packing (token_num_per_sub_group = 2)
{
const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
const int64_t sub_group_idx = block_offset / token_num_per_sub_group;
const int64_t sub_group_offset = block_offset % token_num_per_sub_group;
uint8_t* vdst =
value_cache_ptr + block_idx * vc_stride0 + h * vc_stride1 +
sub_group_idx * token_num_per_sub_group * head_elems_per_group +
sub_group_offset;
for (int64_t i = 0; i < group_num; ++i) {
for (int64_t j = 0; j < head_elems_per_group; ++j)
vdst[j * token_num_per_sub_group] =
quant_fn(static_cast<float>(vsrc[j]), v_inv);
vsrc += head_elems_per_group;
vdst += group_size;
}
}
}
}
}
// ---------------------------------------------------------------------------
// FP8 E5M2 scalar helpers
// ---------------------------------------------------------------------------
// Reference scalar dequant — used to verify vectorized AMX dequant.
// FP8 E5M2: s[7] e[6:2] m[1:0], exponent bias = 15 (same as FP16).
// Byte b → FP16 bits = b << 8 (no bias correction needed).
inline float fp8e5m2_to_float_scalar(uint8_t b, float scale) noexcept {
const uint8_t exp_bits = (b >> 2) & 0x1F;
const uint8_t mant_bits = b & 0x03;
// NaN: exp=11111, mant!=00
if (exp_bits == 0x1F && mant_bits != 0)
return std::numeric_limits<float>::quiet_NaN();
const uint32_t sign = static_cast<uint32_t>(b & 0x80) << 24;
if (exp_bits == 0x1F)
return sign ? -std::numeric_limits<float>::infinity()
: std::numeric_limits<float>::infinity();
if (exp_bits == 0) { // subnormal: (-1)^s * 2^-14 * mant/4
if (mant_bits == 0) return 0.0f;
float v = mant_bits * 0x1p-16f;
return (sign ? -v : v) * scale;
}
// Normal: FP32 exp = exp5 - 15 + 127, mantissa top 2 bits
uint32_t fp32_bits = sign |
((static_cast<uint32_t>(exp_bits) - 15 + 127) << 23) |
(static_cast<uint32_t>(mant_bits) << 21);
float val = *reinterpret_cast<const f32_alias_t*>(&fp32_bits);
return val * scale;
}
inline uint8_t float_to_fp8e5m2_scalar(float v, float inv_scale) noexcept {
v *= inv_scale;
constexpr float fp8_e5m2_max = 57344.0f;
v = std::max(-fp8_e5m2_max, std::min(fp8_e5m2_max, v));
if (v == 0.0f) return 0;
uint32_t bits = *reinterpret_cast<const u32_alias_t*>(&v);
const uint8_t sign = static_cast<uint8_t>((bits >> 24) & 0x80);
const int32_t exp_fp32 = static_cast<int32_t>((bits >> 23) & 0xFF) - 127;
const uint8_t mant2 = static_cast<uint8_t>((bits >> 21) & 0x03);
if (exp_fp32 < -14) { // subnormal in E5M2
const int shift = -14 - exp_fp32;
if (shift + 21 >= 32)
return sign; // underflow: too small for E5M2 subnormal
const uint32_t m = (0x800000u | (bits & 0x7FFFFFu)) >> (shift + 21);
return sign | static_cast<uint8_t>(std::min<uint32_t>(m, 3u));
}
const uint8_t exp5 = static_cast<uint8_t>(exp_fp32 + 15);
return sign | (exp5 << 2) | mant2;
}
// ---------------------------------------------------------------------------
// Select the FP8 quant function at compile time based on kv_cache_t.
// ---------------------------------------------------------------------------
template <typename kv_cache_t>
constexpr auto select_fp8_quant_fn() {
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>)
return float_to_fp8e5m2_scalar;
else
return float_to_fp8e4m3_scalar;
}
// ---------------------------------------------------------------------------
// VEC reshape impl — parameterised on the quantisation function.
// Writes key (column-major) and value (row-major) into uint8 FP8 KV cache.
// The pragma omp must live outside VLLM_DISPATCH_FLOATING_TYPES because
// #pragma cannot appear inside variadic macro arguments.
// ---------------------------------------------------------------------------
template <typename scalar_t, uint8_t (*quant_fn)(float, float)>
inline void reshape_and_cache_fp8_vec_impl(
const scalar_t* key_ptr, const scalar_t* value_ptr, uint8_t* key_cache_ptr,
uint8_t* value_cache_ptr, const int64_t* slot_ptr, int64_t token_num,
int64_t head_num, int64_t head_dim, int64_t block_size, int64_t k_stride0,
int64_t k_stride1, int64_t v_stride0, int64_t v_stride1, int64_t kc_stride0,
int64_t kc_stride1, int64_t vc_stride0, int64_t vc_stride1, float k_inv,
float v_inv) {
#pragma omp parallel for collapse(2) schedule(static)
for (int64_t tok = 0; tok < token_num; ++tok) {
for (int64_t h = 0; h < head_num; ++h) {
const int64_t slot = slot_ptr[tok];
if (slot < 0) continue;
const int64_t block_idx = slot / block_size;
const int64_t block_offset = slot % block_size;
// Key layout: column-major within block
const scalar_t* ksrc = key_ptr + tok * k_stride0 + h * k_stride1;
uint8_t* kdst = key_cache_ptr + block_idx * kc_stride0 + h * kc_stride1 +
block_offset;
for (int64_t i = 0; i < head_dim; ++i)
kdst[i * block_size] = quant_fn(static_cast<float>(ksrc[i]), k_inv);
// Value layout: row-major within block (contiguous head_dim bytes)
const scalar_t* vsrc = value_ptr + tok * v_stride0 + h * v_stride1;
uint8_t* vdst = value_cache_ptr + block_idx * vc_stride0 +
h * vc_stride1 + block_offset * head_dim;
for (int64_t i = 0; i < head_dim; ++i)
vdst[i] = quant_fn(static_cast<float>(vsrc[i]), v_inv);
}
}
}
+12 -63
View File
@@ -12,24 +12,10 @@
#include "cpu/utils.hpp"
namespace cpu_attention {
enum class ISA { AMX, VEC, VEC16, NEON, VXE, VSX };
enum class ISA { AMX, VEC, VEC16, NEON, VXE };
// Mirrors csrc/attention/dtype_fp8.cuh Fp8KVCacheDataType exactly.
enum class Fp8KVCacheDataType {
kAuto = 0,
kFp8E4M3 = 1,
kFp8E5M2 = 2,
};
struct AttentionInput;
template <ISA isa, typename scalar_t, int64_t head_dim,
typename kv_cache_scalar_t = scalar_t>
class AttentionImpl {
public:
void init_from_input(const AttentionInput*) {}
float get_output_v_scale() const noexcept { return 1.0f; }
};
template <ISA isa, typename scalar_t, int64_t head_dim>
class AttentionImpl {};
struct AttentionWorkItemGroup {
int32_t req_id;
@@ -164,9 +150,6 @@ struct AttentionMetadata {
case ISA::VXE:
ss << "VXE, ";
break;
case ISA::VSX:
ss << "VSX, ";
break;
}
ss << "workitem_group_num: " << workitem_group_num
<< ", reduction_item_num: " << reduction_item_num
@@ -797,9 +780,6 @@ struct AttentionInput {
int32_t sliding_window_left;
int32_t sliding_window_right;
float softcap;
// FP8 KV cache scales (used by FP8 attention implementations)
float k_scale_fp8 = 1.0f;
float v_scale_fp8 = 1.0f;
};
#define DEFINE_CPU_ATTENTION_PARAMS \
@@ -1172,11 +1152,7 @@ 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));
@@ -1225,17 +1201,8 @@ class AttentionMainLoop {
vec = vec - max_vec;
// compute exp
#if defined(DEFINE_FAST_EXP)
#ifdef __aarch64__
if constexpr (IsReducedPrecision) {
vec = fast_exp_f16(vec);
} else
#endif
{
vec = fast_exp(vec);
}
#ifdef DEFINE_FAST_EXP
vec = fast_exp(vec);
prob_buffer_vec_t output_vec(vec);
output_vec.save(curr_prob_buffer_iter);
#else
@@ -1291,11 +1258,7 @@ 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);
@@ -1309,15 +1272,8 @@ class AttentionMainLoop {
vec_op::FP32Vec16 vec(curr_logits_buffer_iter);
vec = vec * inv_softcap_scale_vec;
#if defined(DEFINE_FAST_EXP)
#ifdef __aarch64__
if constexpr (IsReducedPrecision) {
vec = fast_exp_f16(vec);
} else
#endif
{
vec = fast_exp(vec);
}
#ifdef DEFINE_FAST_EXP
vec = fast_exp(vec);
vec_op::FP32Vec16 inv_vec = ones_vec / vec;
vec = (vec - inv_vec) / (vec + inv_vec);
#else
@@ -1394,13 +1350,6 @@ class AttentionMainLoop {
}
attention_impl_t attn_impl;
constexpr bool fp8_kv = std::is_same_v<kv_cache_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_t, c10::Float8_e5m2>;
float output_v_scale = 1.0f;
if constexpr (fp8_kv) {
attn_impl.init_from_input(input);
output_v_scale = attn_impl.get_output_v_scale();
}
// general information
const int32_t q_head_num = input->num_heads;
@@ -1780,7 +1729,7 @@ class AttentionMainLoop {
reinterpret_cast<query_t*>(input->output) +
output_buffer_offset,
sum_buffer, actual_q_heads_per_kv,
actual_q_token_num, q_head_num, output_v_scale);
actual_q_token_num, q_head_num);
} else {
const int32_t stride =
actual_q_heads_per_kv * split_kv_q_token_num_threshold;
@@ -1850,7 +1799,7 @@ class AttentionMainLoop {
split_output_buffer,
reinterpret_cast<query_t*>(input->output) + output_buffer_offset,
split_sum_buffer, actual_q_heads_per_kv, curr_output_token_num,
q_head_num, output_v_scale);
q_head_num);
}
}
}
@@ -1974,8 +1923,8 @@ class AttentionMainLoop {
query_t* __restrict__ curr_output_buffer,
float* __restrict__ sum_buffer,
const int32_t q_heads_per_kv,
const int32_t actual_q_token_num, const int32_t q_head_num,
const float v_scale = 1.0f) {
const int32_t actual_q_token_num,
const int32_t q_head_num) {
// final output
using output_vec_t = typename VecTypeTrait<query_t>::vec_t;
@@ -1989,7 +1938,7 @@ class AttentionMainLoop {
curr_partial_output_buffer;
query_t* __restrict__ curr_output_buffer_iter = curr_output_buffer;
for (int32_t head_idx = 0; head_idx < q_heads_per_kv; ++head_idx) {
vec_op::FP32Vec16 inv_sum_scale_vec(v_scale / *curr_sum_buffer);
vec_op::FP32Vec16 inv_sum_scale_vec(1.0 / *curr_sum_buffer);
for (int32_t i = 0; i < group_num_per_head; ++i) {
vec_op::FP32Vec16 vec(curr_partial_output_buffer_iter);
+4 -5
View File
@@ -248,8 +248,8 @@ class TileGemmNeonFMLA {
} // namespace
// this is similar to "ISA::VEC" at the moment
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
template <typename scalar_t, int64_t head_dim>
class AttentionImpl<ISA::NEON, scalar_t, head_dim> {
public:
using query_t = scalar_t;
using q_buffer_t = float;
@@ -343,8 +343,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
const int64_t head_num, const int64_t key_head_num_stride,
const int64_t value_head_num_stride, const int64_t num_blocks,
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
const int64_t block_size, const int64_t block_size_stride,
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
const int64_t block_size, const int64_t block_size_stride) {
#pragma omp parallel for collapse(2)
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -389,7 +388,7 @@ class AttentionImpl<ISA::NEON, scalar_t, head_dim, kv_cache_scalar_t> {
#ifdef ARM_BF16_SUPPORT
// For BF16 on Arm, reuse the BFMMLA kernels with 32-token alignment.
template <int64_t head_dim>
class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim, c10::BFloat16>
class AttentionImpl<ISA::NEON, c10::BFloat16, head_dim>
: public AttentionImplNEONBFMMLA<BLOCK_SIZE_ALIGNMENT, ISA::NEON,
head_dim> {};
#endif
+1 -2
View File
@@ -602,8 +602,7 @@ class AttentionImplNEONBFMMLA {
[[maybe_unused]] const int64_t num_blocks,
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
const int64_t block_size,
[[maybe_unused]] const int64_t block_size_stride,
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
[[maybe_unused]] const int64_t block_size_stride) {
const int64_t k_block_stride = (head_dim / TILE_K) * K_INNER_STRIDE;
const int64_t v_pair_stride =
(block_size / V_TOKENS_PER_ROW_BLOCK) * V_INNER_STRIDE;
+28 -105
View File
@@ -1,37 +1,11 @@
#ifndef CPU_ATTN_VEC_HPP
#define CPU_ATTN_VEC_HPP
#include "cpu_attn_fp8.hpp"
#include "cpu_attn_impl.hpp"
namespace cpu_attention {
namespace {
// Load 32 kv_cache_t elements starting at ptr and return them as two FP32Vec16s
// covering the lower 16 and upper 16 positions.
// For FP8: both halves come from a single BF16Vec32 dequant of 32 bytes.
// For BF16/FP16/FP32: two separate vector loads at ptr and ptr+16.
template <typename kv_cache_t>
FORCE_INLINE std::pair<vec_op::FP32Vec16, vec_op::FP32Vec16> load_b_pair_vec(
const kv_cache_t* ptr) {
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e4m3fn>) {
// BF16 container, but values are in the FP16 exponent range (bias 15 not
// 127).
vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
vec_op::fp8_e4m3_tag{});
return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
} else if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
vec_op::BF16Vec32 bf16_b_reg(reinterpret_cast<const uint8_t*>(ptr),
vec_op::fp8_e5m2_tag{});
return {vec_op::FP32Vec16(bf16_b_reg, 0), vec_op::FP32Vec16(bf16_b_reg, 1)};
} else {
using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
return std::make_pair(vec_op::FP32Vec16(load_vec_t(ptr)),
vec_op::FP32Vec16(load_vec_t(ptr + 16)));
}
}
// 8-2-16 pattern, 8 regs for A, 2 regs for B, 16 regs for C, [8, K] @ [k, 32]
template <typename kv_cache_t>
class TileGemm82 {
@@ -80,7 +54,10 @@ class TileGemm82 {
const int32_t block_size, const int32_t dynamic_k_size,
const bool accum_c) {
static_assert(0 < M && M <= 8);
using load_vec_t = typename VecTypeTrait<kv_cache_t>::vec_t;
kv_cache_t* __restrict__ curr_b_0 = b_tile;
kv_cache_t* __restrict__ curr_b_1 = b_tile + 16;
float* __restrict__ curr_c_0 = c_tile;
float* __restrict__ curr_c_1 = c_tile + 16;
@@ -99,14 +76,16 @@ class TileGemm82 {
}
float* __restrict__ curr_a = a_tile;
kv_cache_t* __restrict__ curr_b = b_tile;
for (int32_t k = 0; k < dynamic_k_size; ++k) {
auto [fp32_b_0_reg, fp32_b_1_reg] = load_b_pair_vec(curr_b);
load_vec_t b_0_reg(curr_b_0);
vec_op::FP32Vec16 fp32_b_0_reg(b_0_reg);
load_vec_t b_1_reg(curr_b_1);
vec_op::FP32Vec16 fp32_b_1_reg(b_1_reg);
float* __restrict__ curr_m_a = curr_a;
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
vec_op::FP32Vec16 a_reg(*curr_m_a);
float v = *curr_m_a;
vec_op::FP32Vec16 a_reg(v);
c_regs[i * 2] = c_regs[i * 2] + a_reg * fp32_b_0_reg;
c_regs[i * 2 + 1] = c_regs[i * 2 + 1] + a_reg * fp32_b_1_reg;
@@ -116,7 +95,8 @@ class TileGemm82 {
// update
curr_a += 1;
curr_b += ldb;
curr_b_0 += ldb;
curr_b_1 += ldb;
}
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
@@ -129,20 +109,15 @@ class TileGemm82 {
});
}
};
} // namespace
// This is a general but naive implementation based on vector instructions
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
static constexpr bool fp8_kv =
std::is_same_v<kv_cache_scalar_t, c10::Float8_e4m3fn> ||
std::is_same_v<kv_cache_scalar_t, c10::Float8_e5m2>;
template <typename scalar_t, int64_t head_dim>
class AttentionImpl<ISA::VEC, scalar_t, head_dim> {
public:
using query_t = scalar_t;
using q_buffer_t = float;
using kv_cache_t = kv_cache_scalar_t;
using kv_cache_t = scalar_t;
using logits_buffer_t = float;
using partial_output_buffer_t = float;
using prob_buffer_t = float;
@@ -154,45 +129,11 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
constexpr static int64_t MaxQHeadNumPerIteration = 8;
constexpr static int64_t HeadDim = head_dim;
constexpr static ISA ISAType = ISA::VEC;
constexpr static bool scale_on_logits = fp8_kv;
float k_scale = 1.0f;
float v_scale = 1.0f;
constexpr static bool scale_on_logits = false; // apply scale on q_buffer
public:
void init_from_input(const AttentionInput* input) {
if constexpr (fp8_kv) {
k_scale = input->k_scale_fp8;
v_scale = input->v_scale_fp8;
}
}
float get_output_v_scale() const noexcept {
if constexpr (fp8_kv) {
// VEC dequant unpacks FP8 into a pseudo-FP16 layout (exponent bias 15).
// E4M3 (bias=7) needs correction 2^(15-7) = 2^8; E5M2 bias matches FP16
// so no correction.
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
return v_scale;
} else {
return v_scale * 0x1p8f;
}
}
return 1.0f;
}
template <template <typename tile_gemm_t> typename attention>
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
if constexpr (fp8_kv) {
// Same bias correction as get_output_v_scale: VEC FP8→pseudo-FP16 dequant
// uses bias 15; E4M3 (bias=7) needs ×2^8, E5M2 (bias=15) needs no
// correction.
if constexpr (std::is_same_v<kv_cache_t, c10::Float8_e5m2>) {
scale *= k_scale;
} else {
scale *= k_scale * 0x1p8f;
}
}
attention<TileGemm82<kv_cache_t>> attention_iteration;
attention_iteration(CPU_ATTENTION_PARAMS);
}
@@ -220,19 +161,17 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
// row-major
}
// Copy q to q_buffer and cast it to fp32.
// FP8: QK scale is folded into execute_attention; copy Q unscaled here.
void copy_q_heads_tile(scalar_t* __restrict__ src,
float* __restrict__ q_buffer, const int32_t q_num,
const int32_t q_heads_per_kv,
const int64_t q_num_stride,
const int64_t q_head_stride, float scale) {
// Copy q to q_buffer and cast it to fp32
static void copy_q_heads_tile(
scalar_t* __restrict__ src, // [q_num, q_heads_per_kv, head_size]
float* __restrict__ q_buffer, const int32_t q_num,
const int32_t q_heads_per_kv, const int64_t q_num_stride,
const int64_t q_head_stride, float scale) {
static_assert(head_dim % 16 == 0);
constexpr int32_t unroll_size = head_dim / 16;
using load_vec_t = typename VecTypeTrait<scalar_t>::vec_t;
const float effective_scale = fp8_kv ? 1.0f : scale;
vec_op::FP32Vec16 scale_vec(effective_scale);
vec_op::FP32Vec16 scale_vec(scale);
for (int32_t q_num_idx = 0; q_num_idx < q_num; ++q_num_idx) {
for (int32_t q_head_idx = 0; q_head_idx < q_heads_per_kv; ++q_head_idx) {
scalar_t* __restrict__ curr_q =
@@ -257,26 +196,13 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
// reshape K as column-major and V as row-major
static void reshape_and_cache(
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
kv_cache_t* __restrict__ key_cache, kv_cache_t* __restrict__ value_cache,
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
const int64_t head_num, const int64_t key_head_num_stride,
const int64_t value_head_num_stride, const int64_t num_blocks,
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
const int64_t block_size, const int64_t block_size_stride,
const float k_inv = 0.0f, const float v_inv = 0.0f) {
if constexpr (fp8_kv) {
constexpr auto qfn = select_fp8_quant_fn<kv_cache_t>();
reshape_and_cache_fp8_vec_impl<scalar_t, qfn>(
key, value, reinterpret_cast<uint8_t*>(key_cache),
reinterpret_cast<uint8_t*>(value_cache), slot_mapping, token_num,
head_num, head_dim, block_size, key_token_num_stride,
key_head_num_stride, value_token_num_stride, value_head_num_stride,
num_blocks_stride, cache_head_num_stride, num_blocks_stride,
cache_head_num_stride, k_inv, v_inv);
return;
}
const int64_t block_size, const int64_t block_size_stride) {
#pragma omp parallel for collapse(2)
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
@@ -294,9 +220,8 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
token_idx * key_token_num_stride +
head_idx * key_head_num_stride;
scalar_t* key_cache_start_ptr =
reinterpret_cast<scalar_t*>(key_cache) +
block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
block_offset;
key_cache + block_idx * num_blocks_stride +
head_idx * cache_head_num_stride + block_offset;
#pragma GCC unroll 8
for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
@@ -309,9 +234,8 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
token_idx * value_token_num_stride +
head_idx * value_head_num_stride;
scalar_t* value_cache_start_ptr =
reinterpret_cast<scalar_t*>(value_cache) +
block_idx * num_blocks_stride + head_idx * cache_head_num_stride +
block_offset * head_dim;
value_cache + block_idx * num_blocks_stride +
head_idx * cache_head_num_stride + block_offset * head_dim;
std::memcpy(value_cache_start_ptr, value_start_ptr,
sizeof(scalar_t) * head_dim);
}
@@ -319,7 +243,6 @@ class AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
}
}
};
} // namespace cpu_attention
#endif
+3 -3
View File
@@ -116,9 +116,9 @@ class TileGemm161 {
} // namespace
// This is a general but naive implementation based on vector instructions
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
class AttentionImpl<ISA::VEC16, scalar_t, head_dim, kv_cache_scalar_t>
: public AttentionImpl<ISA::VEC, scalar_t, head_dim, kv_cache_scalar_t> {
template <typename scalar_t, int64_t head_dim>
class AttentionImpl<ISA::VEC16, scalar_t, head_dim>
: public AttentionImpl<ISA::VEC, scalar_t, head_dim> {
public:
using query_t = scalar_t;
using q_buffer_t = float;
-359
View File
@@ -1,359 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#ifndef CPU_ATTN_VSX_HPP
#define CPU_ATTN_VSX_HPP
#include "cpu_attn_impl.hpp"
#include <altivec.h>
#include <type_traits>
namespace cpu_attention {
namespace {
// ppc64le Vector = 16 bytes (128 bits)
#define BLOCK_SIZE_ALIGNMENT 32
#define HEAD_SIZE_ALIGNMENT 32
#define MAX_Q_HEAD_NUM_PER_ITER 16
template <typename kv_cache_t>
FORCE_INLINE void load_row8_B_as_f32(const kv_cache_t* p, __vector float& b0,
__vector float& b1);
// [1] Float Specialization
template <>
FORCE_INLINE void load_row8_B_as_f32<float>(const float* p, __vector float& b0,
__vector float& b1) {
b0 = vec_xl(0, const_cast<float*>(p));
b1 = vec_xl(0, const_cast<float*>(p + 4));
}
// [2] BFloat16 Specialization (Little Endian ppc64le)
// On ppc64le (LE): BF16 bits should land in the HIGH 16 bits of each float32.
// Byte layout of float32 on LE: [byte0(LSB), byte1, byte2, byte3(MSB)]
// We need BF16 in bytes2-3 (high half) with bytes0-1 zeroed.
// vec_mergeh on LE interleaves elements 0..3: result_i = {a[i], b[i]}
// So vec_mergeh(zeros_u16, raw_u16) gives for each uint16 pair:
// uint16[2i] = zeros[i] -> low 16 bits of uint32 -> zeroed mantissa LSBs
// uint16[2i+1] = raw[i] -> high 16 bits of uint32 -> BF16 bits
// Cast to float32 gives exactly (bf16_bits << 16) per element.
template <>
FORCE_INLINE void load_row8_B_as_f32<c10::BFloat16>(const c10::BFloat16* p,
__vector float& b0,
__vector float& b1) {
__vector unsigned short raw = vec_xl(
0, reinterpret_cast<unsigned short*>(const_cast<c10::BFloat16*>(p)));
__vector unsigned short zeros = vec_splat_u16(0);
// LE: zeros in low 16 bits, raw in high 16 bits → bf16 << 16 == float32
b0 = (__vector float)vec_mergeh(zeros, raw);
b1 = (__vector float)vec_mergel(zeros, raw);
}
// Note: c10::Half (FP16) is not supported on PowerPC architecture
template <int32_t M, typename kv_cache_t>
FORCE_INLINE void gemm_micro_ppc64le_Mx8_Ku4(
const float* __restrict A, // [M x K]
const kv_cache_t* __restrict B, // [K x 8]
float* __restrict C, // [M x 8]
int64_t lda, int64_t ldb, int64_t ldc, int32_t K, bool accumulate) {
static_assert(1 <= M && M <= 8, "M must be in [1,8]");
#define ROWS_APPLY(OP) OP(0) OP(1) OP(2) OP(3) OP(4) OP(5) OP(6) OP(7)
#define IF_M(i) if constexpr (M > (i))
// 1. Define A pointers
#define DECL_A(i) const float* a##i = A + (i) * lda;
ROWS_APPLY(DECL_A)
#undef DECL_A
// 2. Define Accumulators (2 vectors covers 8 columns)
#define DECL_ACC(i) __vector float acc##i##_0, acc##i##_1;
ROWS_APPLY(DECL_ACC)
#undef DECL_ACC
// 3. Initialize Accumulators (Load C or Zero)
#define INIT_ACC(i) \
IF_M(i) { \
if (accumulate) { \
acc##i##_0 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 0)); \
acc##i##_1 = vec_xl(0, const_cast<float*>(C + (i) * ldc + 4)); \
} else { \
acc##i##_0 = vec_splats(0.0f); \
acc##i##_1 = vec_splats(0.0f); \
} \
}
ROWS_APPLY(INIT_ACC)
#undef INIT_ACC
int32_t k = 0;
for (; k + 3 < K; k += 4) {
// Load 4 values of A for each Row M: A[k...k+3]
#define LOAD_A4(i) \
__vector float a##i##v; \
IF_M(i) a##i##v = vec_xl(0, const_cast<float*>(a##i + k));
ROWS_APPLY(LOAD_A4)
#undef LOAD_A4
// FMA for specific lane L of A
// ppc64le: vec_madd(b, vec_splat(a, lane), acc)
#define FMAS_LANE(i, aiv, L) \
IF_M(i) { \
__vector float a_broad = vec_splat(aiv, L); \
acc##i##_0 = vec_madd(b0, a_broad, acc##i##_0); \
acc##i##_1 = vec_madd(b1, a_broad, acc##i##_1); \
}
// Unroll K=0..3
{
__vector float b0, b1;
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 0) * ldb, b0, b1);
#define STEP_K0(i) FMAS_LANE(i, a##i##v, 0)
ROWS_APPLY(STEP_K0)
#undef STEP_K0
}
{
__vector float b0, b1;
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 1) * ldb, b0, b1);
#define STEP_K1(i) FMAS_LANE(i, a##i##v, 1)
ROWS_APPLY(STEP_K1)
#undef STEP_K1
}
{
__vector float b0, b1;
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 2) * ldb, b0, b1);
#define STEP_K2(i) FMAS_LANE(i, a##i##v, 2)
ROWS_APPLY(STEP_K2)
#undef STEP_K2
}
{
__vector float b0, b1;
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)(k + 3) * ldb, b0, b1);
#define STEP_K3(i) FMAS_LANE(i, a##i##v, 3)
ROWS_APPLY(STEP_K3)
#undef STEP_K3
}
#undef FMAS_LANE
}
for (; k < K; ++k) {
__vector float b0, b1;
load_row8_B_as_f32<kv_cache_t>(B + (int64_t)k * ldb, b0, b1);
#define TAIL_ROW(i) \
IF_M(i) { \
__vector float ai = vec_splats(*(a##i + k)); \
acc##i##_0 = vec_madd(b0, ai, acc##i##_0); \
acc##i##_1 = vec_madd(b1, ai, acc##i##_1); \
}
ROWS_APPLY(TAIL_ROW)
#undef TAIL_ROW
}
#define STORE_ROW(i) \
IF_M(i) { \
vec_xst(acc##i##_0, 0, C + (i) * ldc + 0); \
vec_xst(acc##i##_1, 0, C + (i) * ldc + 4); \
}
ROWS_APPLY(STORE_ROW)
#undef STORE_ROW
#undef ROWS_APPLY
#undef IF_M
}
template <int32_t N, typename kv_cache_t>
FORCE_INLINE void gemm_macro_ppc64le_Mx8_Ku4(const float* __restrict A,
const kv_cache_t* __restrict B,
float* __restrict C, int32_t M,
int32_t K, int64_t lda,
int64_t ldb, int64_t ldc,
bool accumulate) {
static_assert(N % 8 == 0, "N must be a multiple of 8");
for (int32_t m = 0; m < M;) {
int32_t mb = (M - m >= 8) ? 8 : (M - m >= 4) ? 4 : (M - m >= 2) ? 2 : 1;
const float* Ab = A + m * lda;
float* Cb = C + m * ldc;
for (int32_t n = 0; n < N; n += 8) {
const kv_cache_t* Bn = B + n;
float* Cn = Cb + n;
switch (mb) {
case 8:
gemm_micro_ppc64le_Mx8_Ku4<8, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
K, accumulate);
break;
case 4:
gemm_micro_ppc64le_Mx8_Ku4<4, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
K, accumulate);
break;
case 2:
gemm_micro_ppc64le_Mx8_Ku4<2, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
K, accumulate);
break;
default:
gemm_micro_ppc64le_Mx8_Ku4<1, kv_cache_t>(Ab, Bn, Cn, lda, ldb, ldc,
K, accumulate);
break;
}
}
m += mb;
}
}
template <typename kv_cache_t>
class TileGemmPPC64 {
public:
template <AttentionGemmPhase phase, int32_t k_size>
FORCE_INLINE static void gemm(const int32_t m_size,
float* __restrict__ a_tile,
kv_cache_t* __restrict__ b_tile,
float* __restrict__ c_tile, const int64_t lda,
const int64_t ldb, const int64_t ldc,
const int32_t block_size,
const int32_t dynamic_k_size,
const bool accum_c) {
if constexpr (phase == AttentionGemmPhase::QK) {
gemm_macro_ppc64le_Mx8_Ku4<BLOCK_SIZE_ALIGNMENT, kv_cache_t>(
a_tile, b_tile, c_tile, m_size, k_size, lda, ldb, ldc, accum_c);
} else {
gemm_macro_ppc64le_Mx8_Ku4<HEAD_SIZE_ALIGNMENT, kv_cache_t>(
a_tile, b_tile, c_tile, m_size, dynamic_k_size, lda, ldb, ldc,
accum_c);
}
}
};
} // namespace
template <typename scalar_t, int64_t head_dim>
class AttentionImpl<ISA::VSX, scalar_t, head_dim> {
public:
using query_t = scalar_t;
using q_buffer_t = float;
using kv_cache_t = scalar_t;
using logits_buffer_t = float;
using partial_output_buffer_t = float;
using prob_buffer_t = float;
constexpr static int64_t BlockSizeAlignment = BLOCK_SIZE_ALIGNMENT;
constexpr static int64_t HeadDimAlignment = HEAD_SIZE_ALIGNMENT;
constexpr static int64_t MaxQHeadNumPerIteration = MAX_Q_HEAD_NUM_PER_ITER;
constexpr static int64_t HeadDim = head_dim;
constexpr static ISA ISAType = ISA::VSX;
constexpr static bool scale_on_logits =
false; // Scale is applied to Q during copy
public:
AttentionImpl() {}
template <template <typename tile_gemm_t> typename attention>
FORCE_INLINE void execute_attention(DEFINE_CPU_ATTENTION_PARAMS) {
attention<TileGemmPPC64<kv_cache_t>> attention_iteration;
attention_iteration(CPU_ATTENTION_PARAMS);
}
// Strides for Memory Layout
constexpr static int64_t k_cache_token_group_stride(
const int32_t block_size) {
return BlockSizeAlignment; // [head_dim, block_size] layout
}
constexpr static int64_t v_cache_token_group_stride(
const int32_t block_size) {
return head_dim * BlockSizeAlignment;
}
constexpr static int64_t v_cache_head_group_stride(const int32_t block_size) {
return HeadDimAlignment;
}
static void copy_q_heads_tile(scalar_t* __restrict__ src,
float* __restrict__ q_buffer,
const int32_t q_num,
const int32_t q_heads_per_kv,
const int64_t q_num_stride,
const int64_t q_head_stride, float scale) {
__vector float scale_vec = vec_splats(scale);
constexpr bool is_bf16 = std::is_same<scalar_t, c10::BFloat16>::value;
for (int32_t i = 0; i < q_num; ++i) {
for (int32_t h = 0; h < q_heads_per_kv; ++h) {
scalar_t* curr_src = src + i * q_num_stride + h * q_head_stride;
float* curr_dst =
q_buffer + i * q_heads_per_kv * head_dim + h * head_dim;
int32_t d = 0;
for (; d <= head_dim - 8; d += 8) {
__vector float v0, v1;
load_row8_B_as_f32<scalar_t>(curr_src + d, v0, v1);
v0 = vec_mul(v0, scale_vec);
v1 = vec_mul(v1, scale_vec);
vec_xst(v0, 0, curr_dst + d);
vec_xst(v1, 0, curr_dst + d + 4);
}
for (; d < head_dim; ++d) {
float val = static_cast<float>(curr_src[d]);
curr_dst[d] = val * scale;
}
}
}
}
static void reshape_and_cache(
const scalar_t* __restrict__ key, const scalar_t* __restrict__ value,
scalar_t* __restrict__ key_cache, scalar_t* __restrict__ value_cache,
const int64_t* __restrict__ slot_mapping, const int64_t token_num,
const int64_t key_token_num_stride, const int64_t value_token_num_stride,
const int64_t head_num, const int64_t key_head_num_stride,
const int64_t value_head_num_stride, const int64_t num_blocks,
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
const int64_t block_size, const int64_t block_size_stride,
const float k_inv = 0.0f, const float v_inv = 0.0f) {
// k_inv and v_inv are unused on VSX: FP8 KV cache is not supported on
// PowerPC. The parameters are present to match the common interface.
#pragma omp parallel for collapse(2)
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
const int64_t pos = slot_mapping[token_idx];
if (pos < 0) continue;
const int64_t block_idx = pos / block_size;
const int64_t block_offset = pos % block_size;
{
const scalar_t* key_src = key + token_idx * key_token_num_stride +
head_idx * key_head_num_stride;
scalar_t* key_dst = key_cache + block_idx * num_blocks_stride +
head_idx * cache_head_num_stride + block_offset;
for (int64_t i = 0, j = 0; i < head_dim; ++i, j += block_size) {
key_dst[j] = key_src[i];
}
}
{
const scalar_t* val_src = value + token_idx * value_token_num_stride +
head_idx * value_head_num_stride;
scalar_t* val_dst = value_cache + block_idx * num_blocks_stride +
head_idx * cache_head_num_stride +
block_offset * head_dim;
std::memcpy(val_dst, val_src, sizeof(scalar_t) * head_dim);
}
}
}
}
};
} // namespace cpu_attention
#undef BLOCK_SIZE_ALIGNMENT
#undef HEAD_SIZE_ALIGNMENT
#undef MAX_Q_HEAD_NUM_PER_ITER
#endif // CPU_ATTN_VSX_HPP
+3 -4
View File
@@ -244,8 +244,8 @@ class TileGemmS390X {
} // namespace
template <typename scalar_t, int64_t head_dim, typename kv_cache_scalar_t>
class AttentionImpl<ISA::VXE, scalar_t, head_dim, kv_cache_scalar_t> {
template <typename scalar_t, int64_t head_dim>
class AttentionImpl<ISA::VXE, scalar_t, head_dim> {
public:
using query_t = scalar_t;
using q_buffer_t = float;
@@ -342,8 +342,7 @@ class AttentionImpl<ISA::VXE, scalar_t, head_dim, kv_cache_scalar_t> {
const int64_t head_num, const int64_t key_head_num_stride,
const int64_t value_head_num_stride, const int64_t num_blocks,
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
const int64_t block_size, const int64_t block_size_stride,
const float /*k_inv*/ = 0.0f, const float /*v_inv*/ = 0.0f) {
const int64_t block_size, const int64_t block_size_stride) {
#pragma omp parallel for collapse(2)
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
-6
View File
@@ -15,9 +15,6 @@ using namespace at::vec;
namespace vec_op {
struct fp8_e4m3_tag {};
struct fp8_e5m2_tag {};
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
@@ -325,9 +322,6 @@ struct BF16Vec32 : public VectorizedRegWrapper<BF16Vec32, 4, c10::BFloat16> {
reg.val[2] = vec8_data.reg.val[0];
reg.val[3] = vec8_data.reg.val[0];
};
explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : Base() {}
explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : Base() {}
};
struct FP32Vec4 : public VectorizedRegWrapper<FP32Vec4, 1, float> {
+12 -26
View File
@@ -15,12 +15,16 @@
#include <torch/all.h>
namespace vec_op {
// 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__)
#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
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -482,18 +486,9 @@ 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)(x, inv_ln2, VEC_ELEM_NUM);
RVVI(__riscv_vfmul_vf_f32, LMUL_256)(reg, 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 =
@@ -711,18 +706,9 @@ 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)(x, inv_ln2, VEC_ELEM_NUM);
RVVI(__riscv_vfmul_vf_f32, LMUL_512)(reg, 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 =
-4
View File
@@ -9,10 +9,6 @@
namespace vec_op {
// FP8 tag types for tag dispatch (see cpu_attn_vec.hpp)
struct fp8_e4m3_tag {};
struct fp8_e5m2_tag {};
// FIXME: FP16 is not fully supported in Torch-CPU
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
-6
View File
@@ -8,9 +8,6 @@
#include <torch/all.h>
namespace vec_op {
struct fp8_e4m3_tag {};
struct fp8_e5m2_tag {};
#define vec_neg(a) (-(a))
#define vec_add(a, b) ((a) + (b))
#define vec_sub(a, b) ((a) - (b))
@@ -244,9 +241,6 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
explicit BF16Vec32(const BF16Vec8& vec8_data)
: reg({vec8_data.reg, vec8_data.reg, vec8_data.reg, vec8_data.reg}) {}
explicit BF16Vec32(const uint8_t*, fp8_e4m3_tag) : reg{} {}
explicit BF16Vec32(const uint8_t*, fp8_e5m2_tag) : reg{} {}
void save(void* ptr) const { *reinterpret_cast<ss16x8x4_t*>(ptr) = reg; }
};
-139
View File
@@ -11,17 +11,6 @@ static_assert(false, "AVX2 must be supported for the current implementation.");
namespace vec_op {
// Tags for FP8 BF16Vec32 constructors (avoid overload collision with
// BF16Vec32(void*)).
// VEC path (FP8 → pseudo-FP16 layout, scale correction applied later):
struct fp8_e4m3_tag {}; // E4M3 → pseudo-FP16; BF16 value = true_E4M3 * 2^-8
struct fp8_e5m2_tag {}; // E5M2 → FP16 bits directly (same exponent bias=15)
// AMX path (FP8 → unscaled BF16, no FP32 round-trip):
// BF16 value = true_E4M3 * 2^-120 (E4M3) or true_E5M2 * 2^-112 (E5M2).
// Exponent rebiasing is folded into k/v scales by the caller.
struct fp8_bf16_e4m3_tag {};
struct fp8_bf16_e5m2_tag {};
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
@@ -187,50 +176,6 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
(__m128i)vec8_data.reg, 2),
(__m128i)vec8_data.reg, 3)) {}
// Decode 32 FP8-E4M3 bytes to pseudo-FP16 layout (stored in the BF16
// register). Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m512i b16 = _mm512_cvtepu8_epi16(b8);
__m512i sign =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
__m512i payload =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 7);
reg = _mm512_or_si512(sign, payload);
}
// Decode 32 FP8-E5M2 bytes to FP16 layout.
// E5M2 and FP16 share the same 5-bit exponent bias (15), so FP8 byte b maps
// directly to FP16 bits by shifting left 8 — no sign/payload reconstruction.
explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
reg = _mm512_slli_epi16(_mm512_cvtepu8_epi16(b8), 8);
}
// Direct FP8-E4M3 → unscaled BF16 for AMX (no FP32 round-trip).
// BF16 value = true_E4M3 * 2^-120; exponent rebiasing folded into k/v scales.
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e4m3_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m512i b16 = _mm512_cvtepu8_epi16(b8);
__m512i sign =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
__m512i payload =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 4);
reg = _mm512_or_si512(sign, payload);
}
// Direct FP8-E5M2 → unscaled BF16 for AMX (no FP32 round-trip).
// BF16 value = true_E5M2 * 2^-112; exponent rebiasing folded into k/v scales.
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e5m2_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m512i b16 = _mm512_cvtepu8_epi16(b8);
__m512i sign =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x80)), 8);
__m512i payload =
_mm512_slli_epi16(_mm512_and_si512(b16, _mm512_set1_epi16(0x7F)), 5);
reg = _mm512_or_si512(sign, payload);
}
void save(void* ptr) const { *reinterpret_cast<__m512i*>(ptr) = reg; }
};
#else
@@ -255,77 +200,6 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
_mm256_castsi128_si256((__m128i)vec8_data.reg),
(__m128i)vec8_data.reg, 1)) {}
// E4M3 decode (AVX2 path) — same bit-layout trick as the AVX512 variant
// above. Result = true_E4M3 * 2^-8; caller applies scale * 2^8.
explicit BF16Vec32(const uint8_t* ptr, fp8_e4m3_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
__m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
__m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
__m256i sign_low = _mm256_slli_epi16(
_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)), 8);
__m256i payload_low = _mm256_slli_epi16(
_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)), 7);
__m256i sign_high = _mm256_slli_epi16(
_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)), 8);
__m256i payload_high = _mm256_slli_epi16(
_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)), 7);
reg_low = _mm256_or_si256(sign_low, payload_low);
reg_high = _mm256_or_si256(sign_high, payload_high);
}
// E5M2 decode (AVX2 path) — b << 8 maps to FP16 bits; see AVX512 variant
// above.
explicit BF16Vec32(const uint8_t* ptr, fp8_e5m2_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
reg_low = _mm256_slli_epi16(_mm256_cvtepu8_epi16(b8_low), 8);
reg_high = _mm256_slli_epi16(_mm256_cvtepu8_epi16(b8_high), 8);
}
// Direct FP8-E4M3 → unscaled BF16 for AMX (AVX2 path, no FP32 round-trip).
// BF16 value = true_E4M3 * 2^-120; exponent rebiasing folded into k/v scales.
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e4m3_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
__m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
__m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
reg_low = _mm256_or_si256(
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)),
8),
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)),
4));
reg_high = _mm256_or_si256(
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)),
8),
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)),
4));
}
// Direct FP8-E5M2 → unscaled BF16 for AMX (AVX2 path, no FP32 round-trip).
// BF16 value = true_E5M2 * 2^-112; exponent rebiasing folded into k/v scales.
explicit BF16Vec32(const uint8_t* ptr, fp8_bf16_e5m2_tag) {
__m256i b8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(ptr));
__m128i b8_low = _mm256_extracti128_si256(b8, 0);
__m128i b8_high = _mm256_extracti128_si256(b8, 1);
__m256i b16_low = _mm256_cvtepu8_epi16(b8_low);
__m256i b16_high = _mm256_cvtepu8_epi16(b8_high);
reg_low = _mm256_or_si256(
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x80)),
8),
_mm256_slli_epi16(_mm256_and_si256(b16_low, _mm256_set1_epi16(0x7F)),
5));
reg_high = _mm256_or_si256(
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x80)),
8),
_mm256_slli_epi16(_mm256_and_si256(b16_high, _mm256_set1_epi16(0x7F)),
5));
}
void save(void* ptr) const {
_mm256_storeu_si256((__m256i*)ptr, reg_low);
_mm256_storeu_si256((__m256i*)ptr + 1, reg_high);
@@ -516,11 +390,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
: reg(_mm512_castsi512_ps(
_mm512_bslli_epi128(_mm512_cvtepu16_epi32(v.reg), 2))) {}
explicit FP32Vec16(const BF16Vec32& v, int upper) {
__m256i v_half_i = _mm512_extracti32x8_epi32(v.reg, upper);
reg = _mm512_cvtph_ps(v_half_i);
}
explicit FP32Vec16(const FP16Vec16& v) : reg(_mm512_cvtph_ps(v.reg)) {}
explicit FP32Vec16(const FP16Vec8& v) : FP32Vec16(FP32Vec8(v)) {}
@@ -625,14 +494,6 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
explicit FP32Vec16(const FP32Vec8& data)
: reg_low(data.reg), reg_high(data.reg) {}
explicit FP32Vec16(const BF16Vec32& v, int upper) {
const __m256i& half = upper ? v.reg_high : v.reg_low;
__m128i lo = _mm256_extractf128_si256(half, 0);
__m128i hi = _mm256_extractf128_si256(half, 1);
reg_low = _mm256_cvtph_ps(lo);
reg_high = _mm256_cvtph_ps(hi);
}
explicit FP32Vec16(const FP16Vec16& v) {
__m128i low = _mm256_extractf128_si256(v.reg, 0);
__m128i high = _mm256_extractf128_si256(v.reg, 1);
+126 -153
View File
@@ -20,98 +20,73 @@ ISA_TYPES = {
"VEC16": 2,
"NEON": 3,
"VXE": 4,
"VSX": 5,
}
# KV cache index: 0 = auto (same as scalar_t), 1 = fp8_e4m3, 2 = fp8_e5m2
KV_CACHE_IDX = {
"auto": 0,
"fp8_e4m3": 1,
"fp8_e5m2": 2,
}
# C++ type for each kv_cache index
KV_CACHE_CPP_TYPES = {
"auto": "scalar_t",
"fp8_e4m3": "c10::Float8_e4m3fn",
"fp8_e5m2": "c10::Float8_e5m2",
}
# ISAs supported for head_dims divisible by 32
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE", "VSX"]
ISA_FOR_32 = ["AMX", "NEON", "VEC", "VEC16", "VXE"]
# ISAs supported for head_dims divisible by 16 only
ISA_FOR_16 = ["VEC16"]
# ISAs that support FP8 KV cache (x86 AVX2/AVX-512 required)
ISA_FOR_FP8 = ["AMX", "VEC"]
def encode_params(head_dim: int, isa_type: str, kv_cache: str = "auto") -> int:
"""Encode head_dim, ISA type, and KV cache type into a single int64_t."""
def encode_params(head_dim: int, isa_type: str) -> int:
"""Encode head_dim and ISA type into a single int64_t."""
isa_val = ISA_TYPES[isa_type]
kv_val = KV_CACHE_IDX[kv_cache]
# Encoding: (head_dim << 16) | (kv_cache_idx << 8) | isa_type
# This allows head_dim up to 2^48 - 1, 256 KV cache types, and 256 ISA types
return (head_dim << 16) | (kv_val << 8) | isa_val
# Encoding: (head_dim << 8) | isa_type
# This allows head_dim up to 2^56 - 1 and 256 ISA types
return (head_dim << 8) | isa_val
def _make_case(
head_dim: int, isa: str, kv_cache: str = "auto", isa_override: str | None = None
) -> str:
"""Generate a single switch case line."""
encoded = encode_params(head_dim, isa, kv_cache)
actual_isa = isa_override if isa_override else isa
cpp_type = KV_CACHE_CPP_TYPES[kv_cache]
attn_impl = (
f"cpu_attention::AttentionImpl<"
f"cpu_attention::ISA::{actual_isa}, \\\n"
f" "
f"scalar_t, head_dim, {cpp_type}>"
)
comment = (
f"head_dim={head_dim}, isa={isa}"
if kv_cache == "auto"
else f"head_dim={head_dim}, isa={isa}, kv_cache={kv_cache}"
)
return (
f""" case {encoded}LL: {{ """
f"""/* {comment} */ \\"""
f"""
constexpr size_t head_dim = {head_dim}; \\"""
f"""
using attn_impl = {attn_impl}; \\"""
f"""
return __VA_ARGS__(); \\"""
f"""
}} \\"""
)
def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False) -> str:
def generate_cases_for_isa_group(isa_list: list[str]) -> str:
"""Generate switch cases for a specific ISA group."""
cases = []
# Non-FP8 cases for head_dims divisible by 32
# Generate cases for head_dims divisible by 32
for head_dim in HEAD_DIMS_32:
for isa in isa_list:
if isa not in ISA_FOR_32:
continue
cases.append(_make_case(head_dim, isa, "auto"))
encoded = encode_params(head_dim, isa)
case_str = (
f""" case {encoded}LL: {{ """
f"""/* head_dim={head_dim}, isa={isa} */ \\"""
f"""
constexpr size_t head_dim = {head_dim}; \\"""
f"""
using attn_impl = cpu_attention::AttentionImpl<"""
f"""cpu_attention::ISA::{isa}, \\"""
f"""
"""
f"""scalar_t, head_dim>; \\"""
f"""
return __VA_ARGS__(); \\"""
f"""
}} \\"""
)
cases.append(case_str)
# Non-FP8 cases for head_dims divisible by 16 only
# Generate cases for head_dims divisible by 16 only
for head_dim in HEAD_DIMS_16:
for isa in isa_list:
cases.append(_make_case(head_dim, isa, "auto", isa_override="VEC16"))
# FP8 cases: only AMX and VEC, only head_dims divisible by 32
if include_fp8:
for fp8_type in ("fp8_e4m3", "fp8_e5m2"):
for head_dim in HEAD_DIMS_32:
for isa in isa_list:
if isa not in ISA_FOR_FP8:
continue
cases.append(_make_case(head_dim, isa, fp8_type))
encoded = encode_params(head_dim, isa)
case_str = (
f""" case {encoded}LL: {{ """
f"""/* head_dim={head_dim}, isa={isa} """
f"""(using VEC16) */ \\"""
f"""
constexpr size_t head_dim = {head_dim}; \\"""
f"""
using attn_impl = cpu_attention::AttentionImpl<"""
f"""cpu_attention::ISA::VEC16, \\"""
f"""
"""
f"""scalar_t, head_dim>; \\"""
f"""
return __VA_ARGS__(); \\"""
f"""
}} \\"""
)
cases.append(case_str)
return "\n".join(cases)
@@ -119,9 +94,8 @@ def generate_cases_for_isa_group(isa_list: list[str], include_fp8: bool = False)
def generate_helper_function() -> str:
"""Generate helper function to encode parameters."""
return """
inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa,
int64_t kv_cache_idx = 0) {
return (head_dim << 16) | (kv_cache_idx << 8) | static_cast<int64_t>(isa);
inline int64_t encode_cpu_attn_params(int64_t head_dim, cpu_attention::ISA isa) {
return (head_dim << 8) | static_cast<int64_t>(isa);
}
"""
@@ -149,94 +123,93 @@ def generate_header_file() -> str:
#include "cpu_attn_vxe.hpp"
#endif
#ifdef __powerpc__
#include "cpu_attn_vsx.hpp"
#endif
"""
header += generate_helper_function()
# Generate dispatch macro with conditional compilation for different ISA sets
header += """
// Dispatch macro using encoded parameters.
// KV_CACHE_IDX: Fp8KVCacheDataType enum value (kAuto=0, kFp8E4M3=1, kFp8E5M2=2).
// FP8 cases (kv_cache_idx != 0) are generated on x86 platforms with AVX2 or
// AVX-512: BF16Vec32 FP8 constructors have both AVX-512 and AVX2 implementations
// in cpu_types_x86.hpp. Non-x86 platforms (#else fallback) have fp8=False.
// Dispatch macro using encoded parameters
"""
def _macro_block(guard: str, isa_list: list[str], fp8: bool) -> str:
"""Return one CPU_ATTN_DISPATCH macro block for a given guard."""
enc = (
" int64_t encoded_params = encode_cpu_attn_params("
"HEAD_DIM, ISA_TYPE, KV_CACHE_IDX); \\"
)
cases = generate_cases_for_isa_group(isa_list, include_fp8=fp8)
tail = (
"\n"
" default: { \\\n"
" TORCH_CHECK(false, "
'"Unsupported CPU attention configuration: head_dim=" + \\\n'
' std::to_string(HEAD_DIM) + " isa=" + \\\n'
" std::to_string(static_cast<int>(ISA_TYPE))"
" + \\\n"
' " kv_cache_idx=" + '
"std::to_string(KV_CACHE_IDX)); \\\n"
" } \\\n"
" } \\\n"
" }()\n\n"
)
return (
f"{guard}\n"
"#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, KV_CACHE_IDX, ...) \\\n"
" [&] { \\\n"
f"{enc}\n"
" switch (encoded_params) { \\\n"
f"{cases}"
f"{tail}"
)
# x86_64 with AMX
header += """#if defined(CPU_CAPABILITY_AMXBF16)
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
[&] { \\
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
switch (encoded_params) { \\
"""
header += generate_cases_for_isa_group(["AMX", "VEC", "VEC16"])
header += """
default: { \\
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
std::to_string(HEAD_DIM) + " isa=" + \\
std::to_string(static_cast<int>(ISA_TYPE))); \\
} \\
} \\
}()
header += _macro_block(
"#if defined(CPU_CAPABILITY_AMXBF16)",
["AMX", "VEC", "VEC16"],
fp8=True,
)
header += _macro_block(
"#elif defined(__aarch64__)",
["NEON", "VEC", "VEC16"],
fp8=False,
)
header += _macro_block(
"#elif defined(__s390x__)",
["VXE", "VEC", "VEC16"],
fp8=False,
)
header += _macro_block(
"#elif defined(__powerpc__)",
["VSX", "VEC", "VEC16"],
fp8=False,
)
header += _macro_block(
"#elif defined(__AVX512F__)",
["VEC", "VEC16"],
fp8=True,
)
header += _macro_block(
"#elif defined(__AVX2__)",
["VEC", "VEC16"],
fp8=False,
)
header += _macro_block(
"#else",
["VEC", "VEC16"],
fp8=False,
)
header += (
"#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / "
"__s390x__ / __powerpc__ */\n\n"
"#endif // CPU_ATTN_DISPATCH_GENERATED_H\n"
)
"""
# ARM64 with NEON
header += """#elif defined(__aarch64__)
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
[&] { \\
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
switch (encoded_params) { \\
"""
header += generate_cases_for_isa_group(["NEON", "VEC", "VEC16"])
header += """
default: { \\
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
std::to_string(HEAD_DIM) + " isa=" + \\
std::to_string(static_cast<int>(ISA_TYPE))); \\
} \\
} \\
}()
"""
# s390x with VXE
header += """#elif defined(__s390x__)
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
[&] { \\
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
switch (encoded_params) { \\
"""
header += generate_cases_for_isa_group(["VXE", "VEC", "VEC16"])
header += """
default: { \\
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
std::to_string(HEAD_DIM) + " isa=" + \\
std::to_string(static_cast<int>(ISA_TYPE))); \\
} \\
} \\
}()
"""
# Fallback: VEC and VEC16 only
header += """#else
#define CPU_ATTN_DISPATCH(HEAD_DIM, ISA_TYPE, ...) \\
[&] { \\
int64_t encoded_params = encode_cpu_attn_params(HEAD_DIM, ISA_TYPE); \\
switch (encoded_params) { \\
"""
header += generate_cases_for_isa_group(["VEC", "VEC16"])
header += """
default: { \\
TORCH_CHECK(false, "Unsupported CPU attention configuration: head_dim=" + \\
std::to_string(HEAD_DIM) + " isa=" + \\
std::to_string(static_cast<int>(ISA_TYPE))); \\
} \\
} \\
}()
#endif /* CPU_CAPABILITY_AMXBF16 / __aarch64__ / __s390x__ */
#endif // CPU_ATTN_DISPATCH_GENERATED_H
"""
return header
+1 -6
View File
@@ -178,12 +178,7 @@ 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,
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");
torch::Tensor& cos_sin_cache, bool is_neox) {
int num_tokens = positions.numel();
int rot_dim = cos_sin_cache.size(1);
int num_heads = query.size(-1) / head_size;
+6 -13
View File
@@ -101,9 +101,7 @@ void cpu_attn_reshape_and_cache(const torch::Tensor& key,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
const torch::Tensor& slot_mapping,
const std::string& isa, const double k_scale,
const double v_scale,
const std::string& kv_cache_dtype);
const std::string& isa);
void cpu_attention_with_kv_cache(
const torch::Tensor& query, const torch::Tensor& key_cache,
@@ -114,8 +112,7 @@ void cpu_attention_with_kv_cache(
const int64_t sliding_window_left, const int64_t sliding_window_right,
const torch::Tensor& block_table, const double softcap,
const torch::Tensor& scheduler_metadata,
const std::optional<torch::Tensor>& s_aux, const double k_scale,
const double v_scale, const std::string& kv_cache_dtype);
const std::optional<torch::Tensor>& s_aux);
// Note: just for avoiding importing errors
void placeholder_op() { TORCH_CHECK(false, "Unimplemented"); }
@@ -266,8 +263,7 @@ 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, int "
"rope_dim_offset=0, bool inverse=False) -> ()");
" Tensor cos_sin_cache, bool is_neox) -> ()");
ops.impl("rotary_embedding", torch::kCPU, &rotary_embedding);
// Quantization
@@ -387,18 +383,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
&get_scheduler_metadata);
ops.def(
"cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
"key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str isa, "
"float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
"()",
"key_cache, Tensor(a3!) value_cache, Tensor slot_mapping, str "
"isa) -> ()",
&cpu_attn_reshape_and_cache);
ops.def(
"cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor "
"value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor "
"seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt "
"sliding_window_left, SymInt sliding_window_right, Tensor block_table, "
"float softcap, Tensor scheduler_metadata, Tensor? s_aux, "
"float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
"()",
"float softcap, Tensor scheduler_metadata, Tensor? s_aux) -> ()",
&cpu_attention_with_kv_cache);
// placeholders
+1 -1
View File
@@ -54,7 +54,7 @@ struct Counter {
};
inline int64_t get_available_l2_size() {
#if defined(__s390x__) || defined(__powerpc__)
#if defined(__s390x__)
static int64_t size = []() {
uint32_t l2_cache_size = 0;
auto caps = at::cpu::get_cpu_capabilities();
+35 -10
View File
@@ -96,14 +96,44 @@ struct enable_sm90_or_later : Kernel {
};
template <typename Kernel>
struct enable_sm100_to_sm120 : Kernel {
struct enable_sm90_only : Kernel {
template <typename... Args>
CUTLASS_DEVICE void operator()(Args&&... args) {
#if defined __CUDA_ARCH__
#if (__CUDA_ARCH__ >= 1000 && __CUDA_ARCH__ < 1200)
#if __CUDA_ARCH__ == 900
Kernel::operator()(std::forward<Args>(args)...);
#else
printf("This kernel only supports sm[100, 120).\n");
printf("This kernel only supports sm90.\n");
asm("trap;");
#endif
#endif
}
};
template <typename Kernel>
struct enable_sm100f_only : Kernel {
template <typename... Args>
CUTLASS_DEVICE void operator()(Args&&... args) {
#if defined __CUDA_ARCH__
#if __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030
Kernel::operator()(std::forward<Args>(args)...);
#else
printf("This kernel only supports sm100f.\n");
asm("trap;");
#endif
#endif
}
};
template <typename Kernel>
struct enable_sm100a_only : Kernel {
template <typename... Args>
CUTLASS_DEVICE void operator()(Args&&... args) {
#if defined __CUDA_ARCH__
#if __CUDA_ARCH__ == 1000
Kernel::operator()(std::forward<Args>(args)...);
#else
printf("This kernel only supports sm100a.\n");
asm("trap;");
#endif
#endif
@@ -118,7 +148,7 @@ struct enable_sm120_only : Kernel {
#if __CUDA_ARCH__ == 1200
Kernel::operator()(std::forward<Args>(args)...);
#else
printf("This kernel only supports sm120a.\n");
printf("This kernel only supports sm120.\n");
asm("trap;");
#endif
#endif
@@ -130,13 +160,8 @@ template <typename Kernel>
struct enable_sm120_family : Kernel {
template <typename... Args>
CUTLASS_DEVICE void operator()(Args&&... args) {
#if defined __CUDA_ARCH__
#if (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
#if defined __CUDA_ARCH__ && (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
Kernel::operator()(std::forward<Args>(args)...);
#else
printf("This kernel only supports sm120f.\n");
asm("trap;");
#endif
#endif
}
};
@@ -1,511 +0,0 @@
/*
* 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>
#ifndef USE_ROCM
#include <cuda_fp8.h>
#else
#include <hip/hip_fp8.h>
#endif
#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
#ifdef USE_ROCM
#define FINAL_MASK 0xffffffffffffffffULL
#else
#define FINAL_MASK 0xffffffffu
#endif
#endif
#ifdef USE_ROCM
// ROCm-compatible FP8 conversion helpers
__device__ __forceinline__ uint8_t rocm_cvt_float_to_fp8_e4m3(float val) {
#if defined(HIP_FP8_TYPE_OCP)
__hip_fp8_e4m3 fp8_val(val);
#else
__hip_fp8_e4m3_fnuz fp8_val(val);
#endif
return reinterpret_cast<uint8_t&>(fp8_val);
}
#endif
namespace vllm {
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);
#ifndef USE_ROCM
__nv_fp8_storage_t s =
__nv_cvt_float_to_fp8(scaled, __NV_SATFINITE, __NV_E4M3);
out_bytes[i] = static_cast<uint8_t>(s);
#else
out_bytes[i] = rocm_cvt_float_to_fp8_e4m3(scaled);
#endif
}
// One 16-byte STG per lane.
*reinterpret_cast<uint4*>(token_fp8_ptr + dim_base) =
*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.
#ifndef USE_ROCM
static int const sm_version = getSMVersion();
// Host-side guard: the device kernel body is compiled as a no-op for
// bf16 on pre-Ampere (sm_70/sm_75) because _typeConvert<BFloat16> is
// 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);
#else
// ROCm: use standard kernel launch syntax (no PDL/stream serialization)
// clang-format off
fusedDeepseekV4QNormRopeKVRopeQuantInsertKernel<scalar_t_in>
<<<grid, kBlockSize, 0, stream>>>(
q_inout, kv_in, k_cache, slot_mapping, position_ids, cos_sin_cache,
eps, num_tokens_full, num_tokens_insert, num_heads_q,
cache_block_size, kv_block_stride);
#endif
}
} // namespace deepseek_v4_fused_ops
} // 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);
});
}
+7 -15
View File
@@ -77,8 +77,7 @@ __global__ void rms_norm_kernel(
#pragma unroll
for (int j = 0; j < VEC_SIZE; j++) {
float x = static_cast<float>(src1.val[j]);
float w = static_cast<float>(src2.val[j]);
dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
dst.val[j] = ((scalar_t)(x * s_variance)) * src2.val[j];
}
v_out[i] = dst;
}
@@ -135,17 +134,10 @@ 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> 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;
_f16Vec<scalar_t, width> temp = residual_v[id];
temp *= s_variance;
temp *= weight_v[idx];
input_v[strided_id] = temp;
}
}
@@ -182,8 +174,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];
float w = (float)weight[idx];
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
input[blockIdx.x * input_stride + idx] =
((scalar_t)(x * s_variance)) * weight[idx];
}
}
+10 -28
View File
@@ -65,16 +65,9 @@ __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 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);
float const out_norm = ((scalar_t)(x * s_variance)) * src2.val[j];
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
scale_inv);
scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
}
}
}
@@ -134,21 +127,13 @@ 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> res = residual_v[id];
_f16Vec<scalar_t, width> w = weight_v[idx];
using Converter = _typeConvert<scalar_t>;
using HipT = typename Converter::hip_type;
_f16Vec<scalar_t, width> temp = residual_v[id];
temp *= s_variance;
temp *= weight_v[idx];
#pragma unroll
for (int i = 0; i < width; ++i) {
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);
out[id * width + i] =
scaled_fp8_conversion<true, fp8_type>(float(temp.data[i]), scale_inv);
}
}
}
@@ -191,12 +176,9 @@ 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 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);
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);
}
}
@@ -277,9 +277,7 @@ 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);
// 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) {
if (n_experts >= 4) {
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
<<<grid, block, shared_mem_size, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
@@ -301,9 +299,7 @@ void quant_impl(void* output, void* output_scale, void* input,
n_experts);
}
} else {
// 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) {
if (n_experts >= 16) {
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
<<<grid, block, 0, stream>>>(
m_topk, k, reinterpret_cast<T*>(input),
@@ -141,7 +141,7 @@ struct cutlass_3x_gemm_sm100 {
sizeof(typename CollectiveEpilogue::SharedStorage))>,
KernelSchedule>::CollectiveOp;
using GemmKernel = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
using GemmKernel = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
};
@@ -125,7 +125,7 @@ struct cutlass_3x_gemm_fp8_blockwise {
MainloopScheduler
>::CollectiveOp>;
using KernelType = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
using KernelType = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue>>;
struct GemmKernel : public KernelType {};
@@ -92,7 +92,7 @@ struct cutlass_3x_gemm_sm100_fp8 {
// -----------------------------------------------------------
// Kernel definition
// -----------------------------------------------------------
using GemmKernel = enable_sm100_to_sm120<cutlass::gemm::kernel::GemmUniversal<
using GemmKernel = enable_sm100f_only<cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
};
@@ -236,41 +236,17 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
#undef LAUNCH_KERNEL
}
// Register-resident fast path for group_size==128.
//
// Each thread holds 16 source elements (32 B = uint4 x 2) in registers across
// the absmax reduce -> scale compute -> quantize pipeline. No shared memory.
// UE8M0 scale extracted via bit math (bit-exact with exp2f(ceilf(log2f))).
//
// Loads two contiguous uint4s (16 B + 16 B = 32 B) per thread; on Blackwell
// nvcc fuses these into a single 256-bit LDG.E.256.
//
// Constraints: GROUP_SIZE % (THREADS_PER_GROUP * VEC_SIZE) == 0; for
// THREADS_PER_GROUP=8 and bf16/fp16 (VEC_SIZE=16), this means GROUP_SIZE=128.
template <typename T, typename DST_DTYPE, int GROUP_SIZE>
__global__ void per_token_group_quant_8bit_packed_register_kernel(
template <typename T, typename DST_DTYPE>
__global__ void per_token_group_quant_8bit_packed_kernel(
const T* __restrict__ input, void* __restrict__ output_q,
unsigned int* __restrict__ output_s_packed, const int64_t num_groups_padded,
const int groups_per_block, const int padded_groups_per_row,
const int groups_per_row, const int mn, const int output_q_mn_extent,
const int tma_aligned_mn, const int64_t num_scale_elems, const float eps,
unsigned int* __restrict__ output_s_packed, const int group_size,
const int num_groups_padded, const int groups_per_block,
const int padded_groups_per_row, const int groups_per_row, const int mn,
const int tma_aligned_mn, const int num_scale_elems, const float eps,
const float min_8bit, const float max_8bit) {
static_assert(GROUP_SIZE == 128, "fast path supports GROUP_SIZE==128");
constexpr int THREADS_PER_GROUP = 8;
constexpr int VEC_SIZE = 32 / sizeof(T); // 16 for bf16/fp16
static_assert(GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE,
"GROUP_SIZE must equal THREADS_PER_GROUP * VEC_SIZE");
// Each group's 8 threads must live in a single warp octet so the
// 0xffu << (threadIdx.x & 24u) shuffle mask selects exactly the lanes
// that share a group. Requires 32 % THREADS_PER_GROUP == 0 and the host
// to launch num_threads as a multiple of THREADS_PER_GROUP (which it does
// via num_threads = groups_per_block * THREADS_PER_GROUP).
static_assert(32 % THREADS_PER_GROUP == 0,
"THREADS_PER_GROUP must divide warp size for the shuffle "
"mask to be valid");
const int local_group_id = threadIdx.x / THREADS_PER_GROUP;
const int lane_id = threadIdx.x % THREADS_PER_GROUP;
const int threads_per_group = 16;
const int64_t local_group_id = threadIdx.x / threads_per_group;
const int lane_id = threadIdx.x % threads_per_group;
const int64_t block_group_id = blockIdx.x * groups_per_block;
const int64_t global_group_id = block_group_id + local_group_id;
@@ -278,207 +254,141 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(
return;
}
// map flat group id to 2D indices (mn_idx, sf_k_idx)
const int sf_k_idx =
static_cast<int>(global_group_id % padded_groups_per_row);
const int mn_idx = static_cast<int>(global_group_id / padded_groups_per_row);
// whether it is a valid group (not padding)
const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row);
// Load 16 input elements (32 B) into registers as two adjacent uint4
// loads. nvcc keeps these as 2x LDG.E.128 on sm_100; the per-thread cost
// is dominated by HBM bandwidth at large MN, so a fused 256-bit load via
// inline PTX gave no measurable speedup.
// alignas(16) is required so the uint4* reinterpret_cast below is
// well-defined for T == bf16/fp16 (default alignof is 2).
alignas(16) T regs[VEC_SIZE];
float local_absmax = eps;
// shared memory to cache each group's data to avoid double DRAM reads.
extern __shared__ __align__(16) char smem_raw[];
T* smem = reinterpret_cast<T*>(smem_raw);
T* smem_group = smem + local_group_id * group_size;
// compute scale for valid groups
float y_s = 0.f;
if (is_valid_group) {
const T* group_input =
input + static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
uint4* dst = reinterpret_cast<uint4*>(&regs[0]);
const uint4* src = reinterpret_cast<const uint4*>(group_input);
dst[0] = src[0];
dst[1] = src[1];
#pragma unroll
for (int i = 0; i < VEC_SIZE; ++i) {
float v = fabsf(static_cast<float>(regs[i]));
local_absmax = fmaxf(local_absmax, v);
}
input + static_cast<int64_t>(mn_idx) * groups_per_row * group_size +
sf_k_idx * group_size;
y_s = ComputeGroupScale<T, true>(group_input, smem_group, group_size,
lane_id, threads_per_group, eps, max_8bit);
}
// 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the
// 8 lanes within the warp that share a group.
unsigned mask = 0xffu << (threadIdx.x & 24u);
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4));
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2));
local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1));
float y_s = local_absmax / max_8bit;
y_s = fmaxf(y_s, 1e-10f);
uint32_t bits = __float_as_uint(y_s);
uint32_t exp_bits = (bits >> 23) & 0xffu;
uint32_t mant_bits = bits & 0x7fffffu;
uint8_t exp_byte =
static_cast<uint8_t>(exp_bits + (mant_bits != 0u ? 1u : 0u));
// Lane 0 writes the packed scale byte.
// pack 4 scales into a uint32 exponent
if (lane_id == 0) {
// each uint32 in output_s_packed stores 4 packed scales
const int sf_k_pack_idx = sf_k_idx / 4;
const int pos = sf_k_idx % 4;
const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx;
if (is_valid_group) {
reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = exp_byte;
// reinterpret the UE8M0 scale y_s as IEEE bits, extract the 8-bit
// exponent, and place it into the correct byte of the 32-bit word.
const unsigned int bits = __float_as_uint(y_s);
const uint8_t exponent = static_cast<uint8_t>((bits >> 23u) & 0xffu);
reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = exponent;
} else if (out_idx < num_scale_elems) {
// write zero for padding groups if within bounds of output_s_packed
reinterpret_cast<uint8_t*>(output_s_packed)[out_idx * 4 + pos] = 0;
}
}
// For padded mn rows that fall within output_q's allocated extent, write
// a uint4 of zeros to keep the buffer clean for downstream TMA loads.
// Skip writes for sf_k padding (those positions don't exist in output_q).
if (!is_valid_group) {
if (sf_k_idx < groups_per_row && mn_idx >= mn &&
mn_idx < output_q_mn_extent) {
DST_DTYPE* group_output =
static_cast<DST_DTYPE*>(output_q) +
static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
*reinterpret_cast<uint4*>(group_output) = make_uint4(0, 0, 0, 0);
}
return;
__syncthreads();
if (is_valid_group) {
DST_DTYPE* group_output =
static_cast<DST_DTYPE*>(output_q) +
static_cast<int64_t>(mn_idx) * groups_per_row * group_size +
sf_k_idx * group_size;
QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
threads_per_group, y_s, min_8bit, max_8bit);
}
// Reconstruct y_s as a power-of-2 float and use its reciprocal.
float y_s_q = __uint_as_float(static_cast<uint32_t>(exp_byte) << 23);
float inv_y = 1.0f / y_s_q;
// Quantize and pack into 16 fp8/int8 bytes (= uint4). VEC_SIZE==16 so we
// fill four 32-bit words, four bytes each.
uint32_t packed_lo = 0;
uint32_t packed_lo_hi = 0;
uint32_t packed_hi_lo = 0;
uint32_t packed_hi = 0;
#pragma unroll
for (int i = 0; i < VEC_SIZE; ++i) {
float q =
fminf(fmaxf(static_cast<float>(regs[i]) * inv_y, min_8bit), max_8bit);
DST_DTYPE qb = DST_DTYPE(q);
uint8_t byte = *reinterpret_cast<uint8_t*>(&qb);
const int shift = (i & 3) * 8;
if (i < 4) {
packed_lo |= static_cast<uint32_t>(byte) << shift;
} else if (i < 8) {
packed_lo_hi |= static_cast<uint32_t>(byte) << shift;
} else if (i < 12) {
packed_hi_lo |= static_cast<uint32_t>(byte) << shift;
} else {
packed_hi |= static_cast<uint32_t>(byte) << shift;
}
}
uint4 packed_out =
make_uint4(packed_lo, packed_lo_hi, packed_hi_lo, packed_hi);
DST_DTYPE* group_output =
static_cast<DST_DTYPE*>(output_q) +
static_cast<int64_t>(mn_idx) * groups_per_row * GROUP_SIZE +
sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE;
*reinterpret_cast<uint4*>(group_output) = packed_out;
}
// Public entry point: register-resident packed quant kernel.
// Constraints: group_size == 128 and bf16/fp16 input.
void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
torch::stable::Tensor& output_q,
torch::stable::Tensor& output_s_packed,
int64_t group_size, double eps,
double min_8bit, double max_8bit) {
STD_TORCH_CHECK(group_size == 128,
"per_token_group_quant_8bit_packed only supports "
"group_size==128, got ",
group_size, ".");
const auto in_dtype = input.scalar_type();
STD_TORCH_CHECK(
in_dtype == torch::headeronly::ScalarType::Half ||
in_dtype == torch::headeronly::ScalarType::BFloat16,
"per_token_group_quant_8bit_packed only supports bf16/fp16 input.");
STD_TORCH_CHECK(input.is_contiguous());
STD_TORCH_CHECK(output_q.is_contiguous());
const int64_t k = input.size(-1);
STD_TORCH_CHECK(k % group_size == 0, "input last dim k=", k,
" is not divisible by group_size=", group_size, ".");
STD_TORCH_CHECK(k % group_size == 0, "Last dimension (", k,
") must be divisible by group_size (", group_size, ").");
const int64_t mn = input.numel() / k;
const int64_t groups_per_row = k / group_size;
STD_TORCH_CHECK(output_s_packed.dim() == 2,
"output_s_packed must be 2D, got dim=", output_s_packed.dim(),
".");
const int64_t k_num_packed_sfk = (groups_per_row + 3) / 4;
const int64_t tma_aligned_mn = ((mn + 3) / 4) * 4;
// output_q may be allocated with extra padded mn rows (e.g.,
// (tma_aligned_mn, k)) so the kernel can zero-fill them in-line and the
// caller can use torch.empty instead of torch.zeros. The grid only covers
// up to tma_aligned_mn, so we cap the extent there.
const int64_t output_q_mn_actual = output_q.numel() / k;
STD_TORCH_CHECK(output_q_mn_actual >= mn,
"output_q must have at least mn rows; got ",
output_q_mn_actual, " rows for mn=", mn, ".");
const int64_t output_q_mn_extent =
output_q_mn_actual < tma_aligned_mn ? output_q_mn_actual : tma_aligned_mn;
STD_TORCH_CHECK(
output_s_packed.scalar_type() == torch::headeronly::ScalarType::Int,
"output_s_packed must be int32 for UE8M0-packed scales.");
"output_s_packed must have dtype int32 for UE8M0-packed scales.");
// DeepGEMM expects SFA scales in MN-major form with shape
// [mn, ceil_div(K, 128 * 4)] and TMA-aligned stride on the last
// dimension.
STD_TORCH_CHECK(output_s_packed.size(0) == mn &&
output_s_packed.size(1) == k_num_packed_sfk,
"output_s_packed shape must be [", mn, ", ", k_num_packed_sfk,
"]; got [", output_s_packed.size(0), ", ",
"], but got [", output_s_packed.size(0), ", ",
output_s_packed.size(1), "].");
// Verify column-major TMA-aligned layout
STD_TORCH_CHECK(output_s_packed.stride(0) == 1 &&
output_s_packed.stride(1) == tma_aligned_mn,
"output_s_packed strides must be [1, ", tma_aligned_mn,
"]; got [", output_s_packed.stride(0), ", ",
"output_s_packed must have strides [1, ", tma_aligned_mn,
"], but got [", output_s_packed.stride(0), ", ",
output_s_packed.stride(1), "].");
cudaStream_t stream = get_current_cuda_stream();
constexpr int THREADS_PER_GROUP = 8;
constexpr int THREADS_PER_GROUP = 16;
// Expand the grid to cover MN and K padding so every byte in
// output_s_packed is written (padding bytes get zeroed by the kernel).
const int64_t padded_groups_per_row = k_num_packed_sfk * 4;
const int64_t num_groups_padded = tma_aligned_mn * padded_groups_per_row;
// Number of elements in output_s_packed.
const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn;
const int groups_per_block = GetGroupsPerBlock(num_groups_padded);
auto dst_type = output_q.scalar_type();
const int64_t num_blocks = num_groups_padded / groups_per_block;
const int num_blocks = num_groups_padded / groups_per_block;
const int num_threads = groups_per_block * THREADS_PER_GROUP;
// CUDA caps grid.x at 2^31 - 1; this fits any realistic shape but guard
// against pathological inputs.
STD_TORCH_CHECK(num_blocks <= static_cast<int64_t>(INT32_MAX),
"per_token_group_quant_8bit_packed grid too large: ",
num_blocks, " blocks (max ", INT32_MAX, ").");
#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \
do { \
dim3 grid(static_cast<unsigned int>(num_blocks)); \
dim3 block(num_threads); \
per_token_group_quant_8bit_packed_register_kernel<T, DST_DTYPE, 128> \
<<<grid, block, 0, stream>>>( \
static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
num_groups_padded, groups_per_block, \
static_cast<int>(padded_groups_per_row), \
static_cast<int>(groups_per_row), static_cast<int>(mn), \
static_cast<int>(output_q_mn_extent), \
static_cast<int>(tma_aligned_mn), num_scale_elems, \
static_cast<float>(eps), static_cast<float>(min_8bit), \
static_cast<float>(max_8bit)); \
#define LAUNCH_PACKED_KERNEL(T, DST_DTYPE) \
do { \
dim3 grid(num_blocks); \
dim3 block(num_threads); \
size_t smem_bytes = \
static_cast<size_t>(groups_per_block) * group_size * sizeof(T); \
per_token_group_quant_8bit_packed_kernel<T, DST_DTYPE> \
<<<grid, block, smem_bytes, stream>>>( \
static_cast<const T*>(input.data_ptr()), output_q.data_ptr(), \
reinterpret_cast<unsigned int*>(output_s_packed.data_ptr()), \
static_cast<int>(group_size), static_cast<int>(num_groups_padded), \
groups_per_block, static_cast<int>(padded_groups_per_row), \
static_cast<int>(groups_per_row), static_cast<int>(mn), \
static_cast<int>(tma_aligned_mn), \
static_cast<int>(num_scale_elems), static_cast<float>(eps), \
static_cast<float>(min_8bit), static_cast<float>(max_8bit)); \
} while (0)
VLLM_STABLE_DISPATCH_HALF_TYPES(
input.scalar_type(), "per_token_group_quant_8bit_packed_register", ([&] {
VLLM_STABLE_DISPATCH_FLOATING_TYPES(
input.scalar_type(), "per_token_group_quant_8bit_packed", ([&] {
if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {
LAUNCH_REG_KERNEL(scalar_t, __nv_fp8_e4m3);
LAUNCH_PACKED_KERNEL(scalar_t, __nv_fp8_e4m3);
} else if (dst_type == torch::headeronly::ScalarType::Char) {
LAUNCH_REG_KERNEL(scalar_t, int8_t);
LAUNCH_PACKED_KERNEL(scalar_t, int8_t);
} else {
STD_TORCH_CHECK(
false,
@@ -487,7 +397,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
}
}));
#undef LAUNCH_REG_KERNEL
#undef LAUNCH_PACKED_KERNEL
}
void per_token_group_quant_fp8(const torch::stable::Tensor& input,
@@ -8,13 +8,3 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
torch::stable::Tensor& output_s,
int64_t group_size, double eps, double min_8bit,
double max_8bit, bool scale_ue8m0 = false);
// Public op: register-resident packed quant for the DeepGEMM Blackwell path.
// Restricted to group_size == 128 and bf16/fp16 input; other configurations
// raise STD_TORCH_CHECK. The legacy shared-memory fallback was removed because
// no production caller (deep_gemm_moe / input_quant_fp8) uses other shapes.
void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
torch::stable::Tensor& output_q,
torch::stable::Tensor& output_s_packed,
int64_t group_size, double eps,
double min_8bit, double max_8bit);
+4 -9
View File
@@ -12,15 +12,6 @@ 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,
@@ -67,6 +58,10 @@ void shuffle_rows(const torch::Tensor& input_tensor,
torch::Tensor& output_tensor);
#ifndef USE_ROCM
// cuBLAS bf16 x bf16 -> fp32 router GEMM (fallback for non-SM90 / batch > 16)
torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
torch::Tensor const& weight);
// DeepSeek V3 optimized router GEMM kernel for SM90+
// Computes output = mat_a @ mat_b.T where:
// mat_a: [num_tokens, hidden_dim] in bf16
+52
View File
@@ -0,0 +1,52 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
// bf16 x bf16 -> fp32 router GEMM via cuBLAS.
// Uses CUBLAS_COMPUTE_32F so bf16 operands accumulate into fp32,
// matching TRT-LLM's cuBLAS fallback behaviour in dsv3RouterGemmOp.
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>
#include <cublas_v2.h>
// cuBLAS column-major math for row-major PyTorch tensors:
// weight[N,K]_row lda=K -> cuBLAS sees (K,N) col-major; CUBLAS_OP_T ->
// (N,K) input[M,K]_row ldb=K -> cuBLAS sees (K,M) col-major; CUBLAS_OP_N
// -> (K,M) out[M,N]_row ldc=N -> cuBLAS sees (N,M) col-major (written as
// output^T)
// cuBLAS: C(N,M) = weight(N,K) @ input(K,M) => C^T = output[M,N]
// params: m=N, n=M, k=K, lda=K (weight), ldb=K (input), ldc=N (output)
torch::Tensor router_gemm_bf16_fp32(torch::Tensor const& input,
torch::Tensor const& weight) {
TORCH_CHECK(input.dtype() == torch::kBFloat16,
"router_gemm_bf16_fp32: input must be bfloat16");
TORCH_CHECK(weight.dtype() == torch::kBFloat16,
"router_gemm_bf16_fp32: weight must be bfloat16");
TORCH_CHECK(input.dim() == 2 && weight.dim() == 2,
"router_gemm_bf16_fp32: input and weight must be 2-D");
TORCH_CHECK(input.size(1) == weight.size(1),
"router_gemm_bf16_fp32: inner dimensions must match");
int64_t const M = input.size(0);
int64_t const N = weight.size(0);
int64_t const K = input.size(1);
auto out = torch::empty({M, N}, input.options().dtype(torch::kFloat32));
cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
TORCH_CUDABLAS_CHECK(
cublasSetStream(handle, at::cuda::getCurrentCUDAStream()));
float const alpha = 1.0f;
float const beta = 0.0f;
TORCH_CUDABLAS_CHECK(cublasGemmEx(
handle, CUBLAS_OP_T, CUBLAS_OP_N, static_cast<int>(N),
static_cast<int>(M), static_cast<int>(K), &alpha, weight.data_ptr(),
CUDA_R_16BF, static_cast<int>(K), input.data_ptr(), CUDA_R_16BF,
static_cast<int>(K), &beta, out.data_ptr(), CUDA_R_32F,
static_cast<int>(N), CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT));
return out;
}

Some files were not shown because too many files have changed in this diff Show More