forked from Karylab-cklius/vllm
Compare commits
215
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7d98b72bd3 | ||
|
|
75ccdf3145 | ||
|
|
c6fe94b4d5 | ||
|
|
46f01a50ac | ||
|
|
f00efc5265 | ||
|
|
b0cb1da1bd | ||
|
|
0e36e3bbd1 | ||
|
|
494845e79f | ||
|
|
0416dab275 | ||
|
|
c8db00b16c | ||
|
|
80c7683923 | ||
|
|
638d6e9757 | ||
|
|
1ad84fea86 | ||
|
|
12213c6795 | ||
|
|
10c75477b0 | ||
|
|
ac36a7a1e7 | ||
|
|
521aa80f71 | ||
|
|
a76df87db8 | ||
|
|
a4904ba903 | ||
|
|
f83de6d44c | ||
|
|
239fc73553 | ||
|
|
76bf55240c | ||
|
|
9a698f3255 | ||
|
|
4080263bb2 | ||
|
|
fc5fda105f | ||
|
|
b07ec92faa | ||
|
|
27ffbfde8d | ||
|
|
229e01e9e1 | ||
|
|
191146dba5 | ||
|
|
f3a920a076 | ||
|
|
149daf0d72 | ||
|
|
917fdb5bf7 | ||
|
|
4b594b4aa1 | ||
|
|
7d10a4cfce | ||
|
|
910cc8543a | ||
|
|
431934522b | ||
|
|
61a09532f2 | ||
|
|
3de4b2bf3c | ||
|
|
b44311b6ef | ||
|
|
b0d7875180 | ||
|
|
53c2f20dd9 | ||
|
|
37e370fe93 | ||
|
|
2dc5a72e7e | ||
|
|
c79ff5f918 | ||
|
|
1a659a0c37 | ||
|
|
0f6cf7f628 | ||
|
|
c79ad3ae21 | ||
|
|
61c9ef986a | ||
|
|
d6dbdb9b0d | ||
|
|
06da482fb4 | ||
|
|
2f75e7f712 | ||
|
|
7c21548ce3 | ||
|
|
9df2f91232 | ||
|
|
387189c429 | ||
|
|
75576c63be | ||
|
|
16aca639b7 | ||
|
|
6049424b7e | ||
|
|
060b5f61dc | ||
|
|
ec59c1579f | ||
|
|
1750e443f2 | ||
|
|
ba18929079 | ||
|
|
0500ca6a58 | ||
|
|
a1c15bcb0f | ||
|
|
4809de7317 | ||
|
|
05781e21dd | ||
|
|
85f638a2b8 | ||
|
|
08e5067561 | ||
|
|
a7d00ec051 | ||
|
|
b8fb56d970 | ||
|
|
96a739289e | ||
|
|
60d443f738 | ||
|
|
1dca300653 | ||
|
|
fca252d59e | ||
|
|
33178f9006 | ||
|
|
b2b8f679d0 | ||
|
|
de6ec294ef | ||
|
|
61e10f0116 | ||
|
|
6e96891ba0 | ||
|
|
47f1b47a73 | ||
|
|
5aab491bc9 | ||
|
|
5812e1a66b | ||
|
|
8950394e0a | ||
|
|
7bb49be4d1 | ||
|
|
c67650f04b | ||
|
|
f890e1dbe2 | ||
|
|
040cbf95cc | ||
|
|
5b3762a7f0 | ||
|
|
4d30c510ce | ||
|
|
6700813f86 | ||
|
|
eb44b3aaa4 | ||
|
|
7a98c7a392 | ||
|
|
0d9e60619b | ||
|
|
1134545b6f | ||
|
|
3e0c887511 | ||
|
|
adfbbc1005 | ||
|
|
adc98f04d0 | ||
|
|
8def3cdde2 | ||
|
|
616c9bd0f4 | ||
|
|
8688a06d67 | ||
|
|
f25953cc59 | ||
|
|
d9aa35161d | ||
|
|
6bcda970fd | ||
|
|
ea0e9c8f2e | ||
|
|
94ed0bf4e0 | ||
|
|
1940c8441e | ||
|
|
72d16aee15 | ||
|
|
e78a0c8e59 | ||
|
|
97a98006b0 | ||
|
|
0a684ab0c0 | ||
|
|
0d9210a502 | ||
|
|
1d874867ea | ||
|
|
2e2e626b40 | ||
|
|
af91f4b3e4 | ||
|
|
2396a61108 | ||
|
|
97a668152b | ||
|
|
58b2012aa2 | ||
|
|
b7c20d0cfa | ||
|
|
a2b1f9fc3b | ||
|
|
642076d26c | ||
|
|
5feb3950e5 | ||
|
|
4ec199b66a | ||
|
|
7ca017778f | ||
|
|
fbfe58133d | ||
|
|
9dd62d80ab | ||
|
|
f878367898 | ||
|
|
bd091079cb | ||
|
|
b23bd73f54 | ||
|
|
e2d7adeb64 | ||
|
|
15cb8e140d | ||
|
|
f007cceb42 | ||
|
|
0a5069e4e3 | ||
|
|
8ce53a616e | ||
|
|
ae10e855ab | ||
|
|
530ee36a0d | ||
|
|
d835ad572c | ||
|
|
47d0597ca2 | ||
|
|
818cf61e91 | ||
|
|
c01618fdc8 | ||
|
|
823eaf667d | ||
|
|
f1f1259692 | ||
|
|
df13b5aef5 | ||
|
|
4938d44a3b | ||
|
|
37bf988c2f | ||
|
|
9459fc6471 | ||
|
|
5245c80564 | ||
|
|
9bc266d923 | ||
|
|
5c9f6557d7 | ||
|
|
dcfebf93f4 | ||
|
|
752bd10647 | ||
|
|
2730b657c4 | ||
|
|
1dcbbd9cac | ||
|
|
ace9fda495 | ||
|
|
ef0aa7ca2f | ||
|
|
e6d1310b2a | ||
|
|
ac5f38a0f7 | ||
|
|
b6ff8a2f50 | ||
|
|
9243e0124e | ||
|
|
df362b2d6d | ||
|
|
7c2acd38b7 | ||
|
|
a287eb163f | ||
|
|
e94243893d | ||
|
|
29c0ec4d63 | ||
|
|
c7ce03bcbd | ||
|
|
c233d90aa8 | ||
|
|
d96aee0951 | ||
|
|
c71a583aa9 | ||
|
|
f12b80c6ef | ||
|
|
da64db78b9 | ||
|
|
425c4eafb0 | ||
|
|
02c01f442b | ||
|
|
fae543015c | ||
|
|
c9be3a8aa1 | ||
|
|
41ea2dd44a | ||
|
|
088c0be268 | ||
|
|
fcd2255d16 | ||
|
|
b5433b6f50 | ||
|
|
cc25f028b7 | ||
|
|
c4cd2bd544 | ||
|
|
5784507da4 | ||
|
|
bf578e1abd | ||
|
|
efed8a1e83 | ||
|
|
11d291511a | ||
|
|
877dae9c68 | ||
|
|
c4dd6d78fd | ||
|
|
ce2aecc4dc | ||
|
|
f38f3d11fb | ||
|
|
d4b4562917 | ||
|
|
7b3192523e | ||
|
|
4c6e2e4b30 | ||
|
|
8502958810 | ||
|
|
ce4bdcbda4 | ||
|
|
d5b1ec2684 | ||
|
|
867ff69733 | ||
|
|
109b736b86 | ||
|
|
69d4f5ef63 | ||
|
|
426d48bfa1 | ||
|
|
26c909ed74 | ||
|
|
fb1d8ccaf5 | ||
|
|
9354f22204 | ||
|
|
17fdd42100 | ||
|
|
472d330c21 | ||
|
|
3b6c96a101 | ||
|
|
4d4e04f452 | ||
|
|
67fe73b2b4 | ||
|
+1 |
ee8f36d0b3 | ||
|
+1 |
f3e9497e92 | ||
|
|
fe784ff22e | ||
|
|
b88abb5036 | ||
|
|
67f9046e4a | ||
|
|
f17be06fbe | ||
|
|
2cab53ddee | ||
|
|
ab0a20d151 | ||
|
|
4a394bfcda | ||
|
|
c95c663049 | ||
|
|
ab3c1aedf3 |
@@ -18,6 +18,8 @@ steps:
|
||||
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
- tests/kernels/mamba/test_cpu_short_conv.py
|
||||
- tests/kernels/mamba/test_causal_conv1d.py
|
||||
- tests/kernels/mamba/test_mamba_ssm.py
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
|
||||
@@ -28,7 +30,9 @@ steps:
|
||||
pytest -x -v -s tests/kernels/test_onednn.py
|
||||
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
|
||||
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
|
||||
|
||||
# Note: SDE can't be downloaded from CI host because of AWS WAF
|
||||
# - label: CPU-Compatibility Tests
|
||||
|
||||
+397
-367
@@ -31,8 +31,46 @@ steps:
|
||||
- text: "What is the release version?"
|
||||
key: release-version
|
||||
|
||||
- group: "Build Python wheels"
|
||||
- group: "Build CUDA 13.0 Python wheels"
|
||||
key: "build-wheels"
|
||||
steps:
|
||||
- label: "Build wheel - aarch64 - CUDA 13.0"
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cuda-13-0
|
||||
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 ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build wheel - x86_64 - CUDA 13.0"
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cuda-13-0
|
||||
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 ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- block: "Unblock to build additional Python wheels"
|
||||
depends_on: ~
|
||||
key: block-build-additional-wheels
|
||||
if: build.env("NIGHTLY") != "1"
|
||||
|
||||
- group: "Build additional Python wheels"
|
||||
key: "build-additional-wheels"
|
||||
depends_on: block-build-additional-wheels
|
||||
allow_dependency_failure: true
|
||||
steps:
|
||||
- label: "Build wheel - aarch64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
@@ -48,20 +86,6 @@ steps:
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build wheel - aarch64 - CUDA 13.0"
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cuda-13-0
|
||||
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 ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build wheel - aarch64 - CPU"
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cpu
|
||||
@@ -113,7 +137,7 @@ steps:
|
||||
- 'mv artifacts/reassembled/wheel "artifacts/dist/$$wheel_name"'
|
||||
- "aws sts get-caller-identity"
|
||||
- "VLLM_WHEEL_PLATFORM=macos bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)"'
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
|
||||
plugins:
|
||||
- aws-assume-role-with-web-identity#v1.6.0:
|
||||
role-arn: arn:aws:iam::936637512419:role/vllm-release-macos-wheel-uploader
|
||||
@@ -133,20 +157,6 @@ steps:
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build wheel - x86_64 - CUDA 13.0"
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cuda-13-0
|
||||
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 ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
|
||||
- 'bash .buildkite/scripts/annotate-build-artifact.sh "$$BUILDKITE_LABEL" "s3://vllm-wheels/$$BUILDKITE_COMMIT/$(cd artifacts/dist && echo *.whl)" release-wheels'
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build wheel - x86_64 - CPU"
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cpu
|
||||
@@ -162,12 +172,26 @@ steps:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Generate and upload wheel indices"
|
||||
key: generate-wheel-indices
|
||||
depends_on: "build-wheels"
|
||||
allow_dependency_failure: true
|
||||
if: build.env("NIGHTLY") != "1"
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "UPDATE_VERSION_INDEX=0 bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
|
||||
|
||||
- label: "Regenerate indices with additional wheels"
|
||||
key: generate-additional-wheel-indices
|
||||
depends_on:
|
||||
- build-wheels
|
||||
- build-additional-wheels
|
||||
- generate-wheel-indices
|
||||
allow_dependency_failure: true
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
|
||||
- 'UPDATE_NIGHTLY_INDEX="$${NIGHTLY:-0}" bash .buildkite/scripts/generate-and-upload-nightly-index.sh'
|
||||
|
||||
- block: "Unblock to build release Docker images"
|
||||
depends_on: ~
|
||||
@@ -566,366 +590,370 @@ steps:
|
||||
#
|
||||
# =============================================================================
|
||||
|
||||
# ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
|
||||
- label: ":rocm: Build ROCm Base Image & Wheels"
|
||||
id: build-rocm-base-wheels
|
||||
- group: "Build ROCm Wheel / Image "
|
||||
key: "build-rocm-wheel-image"
|
||||
depends_on: ~
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- |
|
||||
set -euo pipefail
|
||||
steps:
|
||||
# ROCm Job 1: Build ROCm Base Wheels (with S3 caching)
|
||||
- label: ":rocm: Build ROCm Base Image & Wheels"
|
||||
id: build-rocm-base-wheels
|
||||
depends_on: ~
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- |
|
||||
set -euo pipefail
|
||||
|
||||
# Generate cache key
|
||||
CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
|
||||
ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"
|
||||
# Generate cache key
|
||||
CACHE_KEY=$$(.buildkite/scripts/cache-rocm-base-wheels.sh key)
|
||||
ECR_CACHE_TAG="public.ecr.aws/q9t5s3a7/vllm-release-repo:$${CACHE_KEY}-rocm-base"
|
||||
|
||||
echo "========================================"
|
||||
echo "ROCm Base Build Configuration"
|
||||
echo "========================================"
|
||||
echo " CACHE_KEY: $${CACHE_KEY}"
|
||||
echo " ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
|
||||
echo "========================================"
|
||||
|
||||
# Login to ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
IMAGE_EXISTS=false
|
||||
WHEELS_EXIST=false
|
||||
|
||||
# Check ECR for Docker image
|
||||
echo "========================================"
|
||||
echo "ROCm Base Build Configuration"
|
||||
echo "========================================"
|
||||
echo " CACHE_KEY: $${CACHE_KEY}"
|
||||
echo " ECR_CACHE_TAG: $${ECR_CACHE_TAG}"
|
||||
echo "========================================"
|
||||
|
||||
# Login to ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
IMAGE_EXISTS=false
|
||||
WHEELS_EXIST=false
|
||||
|
||||
# Check ECR for Docker image
|
||||
|
||||
if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
|
||||
IMAGE_EXISTS=true
|
||||
echo "ECR image cache HIT"
|
||||
fi
|
||||
|
||||
# Check S3 for wheels
|
||||
WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
|
||||
if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
|
||||
WHEELS_EXIST=true
|
||||
echo "S3 wheels cache HIT"
|
||||
fi
|
||||
if docker manifest inspect "$${ECR_CACHE_TAG}" > /dev/null 2>&1; then
|
||||
IMAGE_EXISTS=true
|
||||
echo "ECR image cache HIT"
|
||||
fi
|
||||
|
||||
# Check S3 for wheels
|
||||
WHEEL_CACHE_STATUS=$(.buildkite/scripts/cache-rocm-base-wheels.sh check)
|
||||
if [ "$${WHEEL_CACHE_STATUS}" = "hit" ]; then
|
||||
WHEELS_EXIST=true
|
||||
echo "S3 wheels cache HIT"
|
||||
fi
|
||||
|
||||
|
||||
# Scenario 1: Both cached (best case)
|
||||
if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
|
||||
echo ""
|
||||
echo "FULL CACHE HIT - Reusing both image and wheels"
|
||||
echo ""
|
||||
|
||||
# Scenario 1: Both cached (best case)
|
||||
if [ "$${IMAGE_EXISTS}" = "true" ] && [ "$${WHEELS_EXIST}" = "true" ]; then
|
||||
echo ""
|
||||
echo "FULL CACHE HIT - Reusing both image and wheels"
|
||||
echo ""
|
||||
|
||||
# Download wheels
|
||||
.buildkite/scripts/cache-rocm-base-wheels.sh download
|
||||
|
||||
# Save ECR tag for downstream jobs
|
||||
buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
|
||||
|
||||
# Scenario 2: Full rebuild needed
|
||||
else
|
||||
echo ""
|
||||
echo " CACHE MISS - Building from scratch..."
|
||||
echo ""
|
||||
|
||||
# Build full base image and push to ECR
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--file docker/Dockerfile.rocm_base \
|
||||
--tag "$${ECR_CACHE_TAG}" \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
|
||||
--build-arg SCCACHE_REGION_NAME=us-west-2 \
|
||||
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
|
||||
--push \
|
||||
.
|
||||
|
||||
# Build wheel extraction stage
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--file docker/Dockerfile.rocm_base \
|
||||
--tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
|
||||
--target debs_wheel_release \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
|
||||
--build-arg SCCACHE_REGION_NAME=us-west-2 \
|
||||
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
|
||||
--load \
|
||||
.
|
||||
|
||||
# Extract and upload wheels
|
||||
mkdir -p artifacts/rocm-base-wheels
|
||||
cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
|
||||
docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
|
||||
docker rm $${cid}
|
||||
|
||||
.buildkite/scripts/cache-rocm-base-wheels.sh upload
|
||||
# Download wheels
|
||||
.buildkite/scripts/cache-rocm-base-wheels.sh download
|
||||
|
||||
# Save ECR tag for downstream jobs
|
||||
buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
|
||||
|
||||
# Scenario 2: Full rebuild needed
|
||||
else
|
||||
echo ""
|
||||
echo " CACHE MISS - Building from scratch..."
|
||||
echo ""
|
||||
|
||||
# Build full base image and push to ECR
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--file docker/Dockerfile.rocm_base \
|
||||
--tag "$${ECR_CACHE_TAG}" \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
|
||||
--build-arg SCCACHE_REGION_NAME=us-west-2 \
|
||||
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
|
||||
--push \
|
||||
.
|
||||
|
||||
# Build wheel extraction stage
|
||||
DOCKER_BUILDKIT=1 docker buildx build \
|
||||
--file docker/Dockerfile.rocm_base \
|
||||
--tag rocm-base-debs:$${BUILDKITE_BUILD_NUMBER} \
|
||||
--target debs_wheel_release \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
|
||||
--build-arg SCCACHE_REGION_NAME=us-west-2 \
|
||||
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
|
||||
--load \
|
||||
.
|
||||
|
||||
# Extract and upload wheels
|
||||
mkdir -p artifacts/rocm-base-wheels
|
||||
cid=$(docker create rocm-base-debs:$${BUILDKITE_BUILD_NUMBER})
|
||||
docker cp $${cid}:/app/debs/. artifacts/rocm-base-wheels/
|
||||
docker rm $${cid}
|
||||
|
||||
.buildkite/scripts/cache-rocm-base-wheels.sh upload
|
||||
|
||||
# Cache base docker image to ECR
|
||||
docker push "$${ECR_CACHE_TAG}"
|
||||
|
||||
buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
|
||||
|
||||
echo ""
|
||||
echo " Build complete - Image and wheels cached"
|
||||
fi
|
||||
# Cache base docker image to ECR
|
||||
docker push "$${ECR_CACHE_TAG}"
|
||||
|
||||
buildkite-agent meta-data set "rocm-base-image-tag" "$${ECR_CACHE_TAG}"
|
||||
|
||||
echo ""
|
||||
echo " Build complete - Image and wheels cached"
|
||||
fi
|
||||
|
||||
artifact_paths:
|
||||
- "artifacts/rocm-base-wheels/*.whl"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
artifact_paths:
|
||||
- "artifacts/rocm-base-wheels/*.whl"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
|
||||
# ROCm Job 2: Build vLLM ROCm Wheel
|
||||
- label: ":python: Build vLLM ROCm Wheel - x86_64"
|
||||
id: build-rocm-vllm-wheel
|
||||
depends_on:
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
# Download artifacts and prepare Docker image
|
||||
- |
|
||||
set -euo pipefail
|
||||
# ROCm Job 2: Build vLLM ROCm Wheel
|
||||
- label: ":python: Build vLLM ROCm Wheel - x86_64"
|
||||
id: build-rocm-vllm-wheel
|
||||
depends_on:
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
# Download artifacts and prepare Docker image
|
||||
- |
|
||||
set -euo pipefail
|
||||
|
||||
# Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
|
||||
# This fixes version detection when tags are moved/force-pushed
|
||||
echo "Fetching latest tags from origin..."
|
||||
git fetch --tags --force origin
|
||||
|
||||
# Log tag information for debugging version detection
|
||||
echo "========================================"
|
||||
echo "Git Tag Verification"
|
||||
echo "========================================"
|
||||
echo "Current HEAD: $(git rev-parse HEAD)"
|
||||
echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
|
||||
echo ""
|
||||
echo "Recent tags (pointing to commits near HEAD):"
|
||||
git tag -l --sort=-creatordate | head -5
|
||||
echo "setuptools_scm version detection:"
|
||||
pip install -q setuptools_scm 2>/dev/null || true
|
||||
python3 -c "import setuptools_scm; print(' Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo " (setuptools_scm not available in this environment)"
|
||||
echo "========================================"
|
||||
# Ensure git tags are up-to-date (Buildkite's default fetch doesn't update tags)
|
||||
# This fixes version detection when tags are moved/force-pushed
|
||||
echo "Fetching latest tags from origin..."
|
||||
git fetch --tags --force origin
|
||||
|
||||
# Log tag information for debugging version detection
|
||||
echo "========================================"
|
||||
echo "Git Tag Verification"
|
||||
echo "========================================"
|
||||
echo "Current HEAD: $(git rev-parse HEAD)"
|
||||
echo "git describe --tags: $(git describe --tags 2>/dev/null || echo 'No tags found')"
|
||||
echo ""
|
||||
echo "Recent tags (pointing to commits near HEAD):"
|
||||
git tag -l --sort=-creatordate | head -5
|
||||
echo "setuptools_scm version detection:"
|
||||
pip install -q setuptools_scm 2>/dev/null || true
|
||||
python3 -c "import setuptools_scm; print(' Detected version:', setuptools_scm.get_version())" 2>/dev/null || echo " (setuptools_scm not available in this environment)"
|
||||
echo "========================================"
|
||||
|
||||
# Download wheel artifacts from current build
|
||||
echo "Downloading wheel artifacts from current build"
|
||||
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
|
||||
# Download wheel artifacts from current build
|
||||
echo "Downloading wheel artifacts from current build"
|
||||
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
|
||||
|
||||
# Get ECR image tag from metadata (set by build-rocm-base-wheels)
|
||||
ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
|
||||
if [ -z "$${ECR_IMAGE_TAG}" ]; then
|
||||
echo "ERROR: rocm-base-image-tag metadata not found"
|
||||
echo "This should have been set by the build-rocm-base-wheels job"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
|
||||
|
||||
# Login to ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
# Pull base Docker image from ECR
|
||||
docker pull "$${ECR_IMAGE_TAG}"
|
||||
|
||||
echo "Loaded base image: $${ECR_IMAGE_TAG}"
|
||||
|
||||
# Prepare base wheels for Docker build context
|
||||
mkdir -p docker/context/base-wheels
|
||||
touch docker/context/base-wheels/.keep
|
||||
cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
|
||||
echo "Base wheels for vLLM build:"
|
||||
ls -lh docker/context/base-wheels/
|
||||
# Get ECR image tag from metadata (set by build-rocm-base-wheels)
|
||||
ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
|
||||
if [ -z "$${ECR_IMAGE_TAG}" ]; then
|
||||
echo "ERROR: rocm-base-image-tag metadata not found"
|
||||
echo "This should have been set by the build-rocm-base-wheels job"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
|
||||
|
||||
# Login to ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
# Pull base Docker image from ECR
|
||||
docker pull "$${ECR_IMAGE_TAG}"
|
||||
|
||||
echo "Loaded base image: $${ECR_IMAGE_TAG}"
|
||||
|
||||
# Prepare base wheels for Docker build context
|
||||
mkdir -p docker/context/base-wheels
|
||||
touch docker/context/base-wheels/.keep
|
||||
cp artifacts/rocm-base-wheels/*.whl docker/context/base-wheels/
|
||||
echo "Base wheels for vLLM build:"
|
||||
ls -lh docker/context/base-wheels/
|
||||
|
||||
echo "========================================"
|
||||
echo "Building vLLM wheel with:"
|
||||
echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
|
||||
echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
|
||||
echo " BASE_IMAGE: $${ECR_IMAGE_TAG}"
|
||||
echo "========================================"
|
||||
echo "========================================"
|
||||
echo "Building vLLM wheel with:"
|
||||
echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
|
||||
echo " BUILDKITE_BRANCH: $${BUILDKITE_BRANCH}"
|
||||
echo " BASE_IMAGE: $${ECR_IMAGE_TAG}"
|
||||
echo "========================================"
|
||||
|
||||
# Build vLLM wheel using local checkout (REMOTE_VLLM=0)
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--file docker/Dockerfile.rocm \
|
||||
--target export_vllm_wheel_release \
|
||||
--output type=local,dest=rocm-dist \
|
||||
--build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
|
||||
--build-arg REMOTE_VLLM=0 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
|
||||
--build-arg SCCACHE_REGION_NAME=us-west-2 \
|
||||
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
|
||||
.
|
||||
echo "Built vLLM wheel:"
|
||||
ls -lh rocm-dist/*.whl
|
||||
# Copy wheel to artifacts directory
|
||||
mkdir -p artifacts/rocm-vllm-wheel
|
||||
cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
|
||||
echo "Final vLLM wheel:"
|
||||
ls -lh artifacts/rocm-vllm-wheel/
|
||||
artifact_paths:
|
||||
- "artifacts/rocm-vllm-wheel/*.whl"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
# Build vLLM wheel using local checkout (REMOTE_VLLM=0)
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--file docker/Dockerfile.rocm \
|
||||
--target export_vllm_wheel_release \
|
||||
--output type=local,dest=rocm-dist \
|
||||
--build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
|
||||
--build-arg REMOTE_VLLM=0 \
|
||||
--build-arg GIT_REPO_CHECK=1 \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
|
||||
--build-arg SCCACHE_REGION_NAME=us-west-2 \
|
||||
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
|
||||
.
|
||||
echo "Built vLLM wheel:"
|
||||
ls -lh rocm-dist/*.whl
|
||||
# Copy wheel to artifacts directory
|
||||
mkdir -p artifacts/rocm-vllm-wheel
|
||||
cp rocm-dist/*.whl artifacts/rocm-vllm-wheel/
|
||||
echo "Final vLLM wheel:"
|
||||
ls -lh artifacts/rocm-vllm-wheel/
|
||||
artifact_paths:
|
||||
- "artifacts/rocm-vllm-wheel/*.whl"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
|
||||
# ROCm Job 3: Upload Wheels to S3
|
||||
- label: ":s3: Upload ROCm Wheels to S3"
|
||||
id: upload-rocm-wheels
|
||||
depends_on:
|
||||
- step: build-rocm-vllm-wheel
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
# Download all wheel artifacts and run upload
|
||||
- |
|
||||
set -euo pipefail
|
||||
# ROCm Job 3: Upload Wheels to S3
|
||||
- label: ":s3: Upload ROCm Wheels to S3"
|
||||
id: upload-rocm-wheels
|
||||
depends_on:
|
||||
- step: build-rocm-vllm-wheel
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
# Download all wheel artifacts and run upload
|
||||
- |
|
||||
set -euo pipefail
|
||||
|
||||
# Download artifacts from current build
|
||||
echo "Downloading artifacts from current build"
|
||||
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
|
||||
buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
|
||||
# Download artifacts from current build
|
||||
echo "Downloading artifacts from current build"
|
||||
buildkite-agent artifact download "artifacts/rocm-base-wheels/*.whl" .
|
||||
buildkite-agent artifact download "artifacts/rocm-vllm-wheel/*.whl" .
|
||||
|
||||
# Run upload script
|
||||
bash .buildkite/scripts/upload-rocm-wheels.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
# # Run upload script
|
||||
bash .buildkite/scripts/upload-rocm-wheels.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
|
||||
# ROCm Job 4: Annotate ROCm Wheel Release
|
||||
- label: ":memo: Annotate ROCm wheel release"
|
||||
id: annotate-rocm-release
|
||||
depends_on:
|
||||
- upload-rocm-wheels
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-rocm-release.sh"
|
||||
env:
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
# ROCm Job 4: Annotate ROCm Wheel Release
|
||||
- label: ":memo: Annotate ROCm wheel release"
|
||||
id: annotate-rocm-release
|
||||
depends_on:
|
||||
- upload-rocm-wheels
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/annotate-rocm-release.sh"
|
||||
env:
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
|
||||
# ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
|
||||
# This is the job to create https://wheels.vllm.ai/rocm/ index allowing
|
||||
# users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
|
||||
- block: "Generate Root Index for ROCm Wheels for Release"
|
||||
key: block-generate-root-index-rocm-wheels
|
||||
depends_on: upload-rocm-wheels
|
||||
# ROCm Job 5: Generate Root Index for ROCm Wheels (for release only)
|
||||
# This is the job to create https://wheels.vllm.ai/rocm/ index allowing
|
||||
# users to install with `uv pip install vllm --extra-index-url https://wheels.vllm.ai/rocm/`
|
||||
- block: "Generate Root Index for ROCm Wheels for Release"
|
||||
key: block-generate-root-index-rocm-wheels
|
||||
depends_on: upload-rocm-wheels
|
||||
|
||||
- label: ":package: Generate Root Index for ROCm Wheels for Release"
|
||||
depends_on: block-generate-root-index-rocm-wheels
|
||||
id: generate-root-index-rocm-wheels
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
|
||||
env:
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
VARIANT: "rocm723"
|
||||
- label: ":package: Generate Root Index for ROCm Wheels for Release"
|
||||
depends_on: block-generate-root-index-rocm-wheels
|
||||
id: generate-root-index-rocm-wheels
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
commands:
|
||||
- "bash tools/vllm-rocm/generate-rocm-wheels-root-index.sh"
|
||||
env:
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
VARIANT: "rocm723"
|
||||
|
||||
# ROCm Job 6: Build ROCm Release Docker Image
|
||||
- label: ":docker: Build release image - x86_64 - ROCm"
|
||||
id: build-rocm-release-image
|
||||
depends_on:
|
||||
- step: block-build-release-images
|
||||
allow_failure: true
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- |
|
||||
set -euo pipefail
|
||||
|
||||
# Login to ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
# Get ECR image tag from metadata (set by build-rocm-base-wheels)
|
||||
ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
|
||||
if [ -z "$${ECR_IMAGE_TAG}" ]; then
|
||||
echo "ERROR: rocm-base-image-tag metadata not found"
|
||||
echo "This should have been set by the build-rocm-base-wheels job"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
|
||||
|
||||
# Pull base Docker image from ECR
|
||||
docker pull "$${ECR_IMAGE_TAG}"
|
||||
|
||||
echo "Loaded base image: $${ECR_IMAGE_TAG}"
|
||||
|
||||
# Pass the base image ECR tag to downstream steps (nightly publish)
|
||||
buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}"
|
||||
|
||||
echo "========================================"
|
||||
echo "Building vLLM ROCm release image with:"
|
||||
echo " BASE_IMAGE: $${ECR_IMAGE_TAG}"
|
||||
echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
|
||||
echo "========================================"
|
||||
|
||||
# Build vLLM ROCm release image using cached base
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
|
||||
--build-arg SCCACHE_REGION_NAME=us-west-2 \
|
||||
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile.rocm .
|
||||
|
||||
# Push to ECR
|
||||
docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
|
||||
# ROCm Job 6: Build ROCm Release Docker Image
|
||||
- label: ":docker: Build release image - x86_64 - ROCm"
|
||||
id: build-rocm-release-image
|
||||
depends_on:
|
||||
- step: block-build-release-images
|
||||
allow_failure: true
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
queue: cpu_queue_release
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- |
|
||||
set -euo pipefail
|
||||
|
||||
# Login to ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | \
|
||||
docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
|
||||
# Get ECR image tag from metadata (set by build-rocm-base-wheels)
|
||||
ECR_IMAGE_TAG="$$(buildkite-agent meta-data get rocm-base-image-tag 2>/dev/null || echo '')"
|
||||
if [ -z "$${ECR_IMAGE_TAG}" ]; then
|
||||
echo "ERROR: rocm-base-image-tag metadata not found"
|
||||
echo "This should have been set by the build-rocm-base-wheels job"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Pulling base Docker image from ECR: $${ECR_IMAGE_TAG}"
|
||||
|
||||
# Pull base Docker image from ECR
|
||||
docker pull "$${ECR_IMAGE_TAG}"
|
||||
|
||||
echo "Loaded base image: $${ECR_IMAGE_TAG}"
|
||||
|
||||
# Pass the base image ECR tag to downstream steps (nightly publish)
|
||||
buildkite-agent meta-data set "rocm-base-ecr-tag" "$${ECR_IMAGE_TAG}"
|
||||
|
||||
echo "========================================"
|
||||
echo "Building vLLM ROCm release image with:"
|
||||
echo " BASE_IMAGE: $${ECR_IMAGE_TAG}"
|
||||
echo " BUILDKITE_COMMIT: $${BUILDKITE_COMMIT}"
|
||||
echo "========================================"
|
||||
|
||||
# Build vLLM ROCm release image using cached base
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg BASE_IMAGE="$${ECR_IMAGE_TAG}" \
|
||||
--build-arg USE_SCCACHE=1 \
|
||||
--build-arg SCCACHE_BUCKET_NAME=vllm-build-sccache \
|
||||
--build-arg SCCACHE_REGION_NAME=us-west-2 \
|
||||
--build-arg SCCACHE_S3_NO_CREDENTIALS=0 \
|
||||
--tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm \
|
||||
--target vllm-openai \
|
||||
--progress plain \
|
||||
-f docker/Dockerfile.rocm .
|
||||
|
||||
# Push to ECR
|
||||
docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm
|
||||
|
||||
echo ""
|
||||
echo " Successfully built and pushed ROCm release image"
|
||||
echo " Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
|
||||
echo ""
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
echo ""
|
||||
echo " Successfully built and pushed ROCm release image"
|
||||
echo " Image: public.ecr.aws/q9t5s3a7/vllm-release-repo:$${BUILDKITE_COMMIT}-rocm"
|
||||
echo ""
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
S3_BUCKET: "vllm-wheels"
|
||||
|
||||
- label: "Publish nightly XPU image to DockerHub"
|
||||
depends_on:
|
||||
- create-manifest-xpu
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
password-env: DOCKERHUB_TOKEN
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
- label: "Publish nightly XPU image to DockerHub"
|
||||
depends_on:
|
||||
- create-manifest-xpu
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/xpu/push-nightly-builds-xpu.sh"
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-xpu"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
password-env: DOCKERHUB_TOKEN
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
|
||||
- label: "Publish nightly ROCm image to DockerHub"
|
||||
depends_on:
|
||||
- build-rocm-release-image
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/push-nightly-builds-rocm.sh"
|
||||
# Clean up old nightly builds (keep only last 14)
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm"
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
password-env: DOCKERHUB_TOKEN
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
- label: "Publish nightly ROCm image to DockerHub"
|
||||
depends_on:
|
||||
- build-rocm-release-image
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
agents:
|
||||
queue: small_cpu_queue_release
|
||||
commands:
|
||||
- "bash .buildkite/scripts/push-nightly-builds-rocm.sh"
|
||||
# Clean up old nightly builds (keep only last 14)
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh nightly- vllm/vllm-openai-rocm"
|
||||
- "bash .buildkite/scripts/cleanup-nightly-builds.sh base-nightly- vllm/vllm-openai-rocm"
|
||||
plugins:
|
||||
- docker-login#v3.0.0:
|
||||
username: vllmbot
|
||||
password-env: DOCKERHUB_TOKEN
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
DOCKERHUB_USERNAME: "vllmbot"
|
||||
|
||||
# =============================================================================
|
||||
# Publish to DockerHub and PyPI (at the end so all builds complete first)
|
||||
@@ -974,6 +1002,8 @@ steps:
|
||||
depends_on:
|
||||
- input-release-version
|
||||
- build-wheels
|
||||
- build-additional-wheels
|
||||
- generate-additional-wheel-indices
|
||||
|
||||
- label: "Upload release wheels to PyPI"
|
||||
depends_on:
|
||||
|
||||
@@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm"
|
||||
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
|
||||
DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
|
||||
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
|
||||
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_rixl build_rocshmem build_deepep mori_base ci_base"
|
||||
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
|
||||
DEFAULT_CI_BASE_METADATA_VERSION="1"
|
||||
IMAGE_EXISTED_BEFORE_BUILD=0
|
||||
|
||||
@@ -1159,8 +1159,8 @@ ci_base_metadata_pairs() {
|
||||
metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
|
||||
metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
|
||||
metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
|
||||
metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
|
||||
metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
|
||||
metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
|
||||
metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
|
||||
@@ -1169,7 +1169,7 @@ ci_base_metadata_pairs() {
|
||||
metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
|
||||
metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
|
||||
metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
|
||||
metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
|
||||
metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
|
||||
|
||||
@@ -1686,7 +1686,7 @@ extract_dependency_pins() {
|
||||
return 0
|
||||
fi
|
||||
|
||||
for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
|
||||
for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
|
||||
if [[ -n "${!var:-}" ]]; then
|
||||
echo "Using provided ${var}: ${!var}"
|
||||
continue
|
||||
@@ -1706,30 +1706,30 @@ extract_dependency_pins() {
|
||||
compute_dependency_cache_keys() {
|
||||
local bake_dir=""
|
||||
local dockerfile_rocm=""
|
||||
local rixl_branch=""
|
||||
local nixl_branch=""
|
||||
local ucx_branch=""
|
||||
local rocshmem_branch=""
|
||||
local deepep_branch=""
|
||||
local rixl_material=""
|
||||
local nixl_material=""
|
||||
local rocshmem_material=""
|
||||
local deepep_material=""
|
||||
|
||||
bake_dir=$(dirname "${VLLM_BAKE_FILE}")
|
||||
dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
|
||||
rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
|
||||
nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
|
||||
ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
|
||||
rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
|
||||
deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
|
||||
|
||||
if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
|
||||
rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
|
||||
RIXL_CACHE_KEY=$(
|
||||
if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
|
||||
nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
|
||||
NIXL_CACHE_KEY=$(
|
||||
compose_dependency_cache_key \
|
||||
"${rixl_branch}-ucx-${ucx_branch}" \
|
||||
"${rixl_material}"
|
||||
"${nixl_branch}-ucx-${ucx_branch}" \
|
||||
"${nixl_material}"
|
||||
)
|
||||
export RIXL_CACHE_KEY
|
||||
echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
|
||||
export NIXL_CACHE_KEY
|
||||
echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
|
||||
fi
|
||||
|
||||
if [[ -n "${rocshmem_branch}" ]]; then
|
||||
@@ -1780,11 +1780,11 @@ dependency_cache_ref_for_target() {
|
||||
local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
|
||||
|
||||
case "${target}" in
|
||||
rixl-rocm-ci)
|
||||
if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
|
||||
elif [[ -n "${RIXL_BRANCH:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
|
||||
nixl-rocm-ci)
|
||||
if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
|
||||
elif [[ -n "${NIXL_BRANCH:-}" ]]; then
|
||||
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
|
||||
fi
|
||||
;;
|
||||
rocshmem-rocm-ci)
|
||||
@@ -1815,7 +1815,7 @@ add_dependency_cache_target() {
|
||||
|
||||
resolve_ci_base_dependency_targets() {
|
||||
local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
|
||||
local rixl_ref=""
|
||||
local nixl_ref=""
|
||||
local rocshmem_ref=""
|
||||
local deepep_ref=""
|
||||
|
||||
@@ -1824,7 +1824,7 @@ resolve_ci_base_dependency_targets() {
|
||||
case "${mode}" in
|
||||
always)
|
||||
echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
|
||||
for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
|
||||
for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
|
||||
if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
|
||||
add_dependency_cache_target "${target}"
|
||||
fi
|
||||
@@ -1844,13 +1844,13 @@ resolve_ci_base_dependency_targets() {
|
||||
;;
|
||||
esac
|
||||
|
||||
if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
|
||||
rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
|
||||
if dependency_cache_ref_exists "${rixl_ref}"; then
|
||||
echo "RIXL dependency cache exists: ${rixl_ref}"
|
||||
if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
|
||||
nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
|
||||
if dependency_cache_ref_exists "${nixl_ref}"; then
|
||||
echo "NIXL dependency cache exists: ${nixl_ref}"
|
||||
else
|
||||
echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
|
||||
add_dependency_cache_target "rixl-rocm-ci"
|
||||
echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
|
||||
add_dependency_cache_target "nixl-rocm-ci"
|
||||
fi
|
||||
fi
|
||||
|
||||
|
||||
@@ -45,8 +45,10 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --curr
|
||||
echo "Uploading indices to $S3_COMMIT_PREFIX"
|
||||
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
|
||||
|
||||
# copy to /nightly/ only if it is on the main branch and not a PR
|
||||
if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]]; then
|
||||
# copy to /nightly/ only when enabled for a main branch build that is not a PR
|
||||
if [[ "${UPDATE_NIGHTLY_INDEX:-1}" == "1" && \
|
||||
"$BUILDKITE_BRANCH" == "main" && \
|
||||
"$BUILDKITE_PULL_REQUEST" == "false" ]]; then
|
||||
echo "Uploading indices to overwrite /nightly/"
|
||||
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/nightly/"
|
||||
fi
|
||||
@@ -67,7 +69,7 @@ pure_version="${version%%+*}"
|
||||
echo "Pure version (without variant): $pure_version"
|
||||
|
||||
# re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
|
||||
if [[ "$version" != *"dev"* ]]; then
|
||||
if [[ "${UPDATE_VERSION_INDEX:-1}" == "1" && "$version" != *"dev"* ]]; then
|
||||
echo "Re-generating indices for /$pure_version/"
|
||||
rm -rf "${INDICES_OUTPUT_DIR:?}"
|
||||
mkdir -p "$INDICES_OUTPUT_DIR"
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
#!/bin/bash
|
||||
set -euox pipefail
|
||||
|
||||
export VLLM_CPU_KVCACHE_SPACE=1
|
||||
export VLLM_CPU_KVCACHE_SPACE=1
|
||||
export VLLM_CPU_CI_ENV=1
|
||||
# Reduce sub-processes for acceleration
|
||||
export TORCH_COMPILE_DISABLE=1
|
||||
# Skip torch.compile via vLLM's --enforce-eager flag (passed below) instead of
|
||||
# TORCH_COMPILE_DISABLE=1, which torch 2.12 no longer treats as a silent no-op
|
||||
# when callers specify fullgraph=True.
|
||||
export VLLM_ENABLE_V1_MULTIPROCESSING=0
|
||||
|
||||
SDE_ARCHIVE="sde-external-10.7.0-2026-02-18-lin.tar.xz"
|
||||
@@ -49,15 +50,15 @@ wait_for_pid_and_check_log() {
|
||||
}
|
||||
|
||||
# Test Sky Lake (AVX512F)
|
||||
./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_0.log 2>&1 &
|
||||
./sde/sde64 -skl -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_0.log 2>&1 &
|
||||
PID_TEST_0=$!
|
||||
|
||||
# Test Cascade Lake (AVX512F + VNNI)
|
||||
./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_1.log 2>&1 &
|
||||
./sde/sde64 -clx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_1.log 2>&1 &
|
||||
PID_TEST_1=$!
|
||||
|
||||
# Test Cooper Lake (AVX512F + VNNI + BF16)
|
||||
./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 > test_2.log 2>&1 &
|
||||
./sde/sde64 -cpx -- python3 examples/basic/offline_inference/generate.py --model facebook/opt-125m --dtype bfloat16 --enforce-eager > test_2.log 2>&1 &
|
||||
PID_TEST_2=$!
|
||||
|
||||
wait_for_pid_and_check_log $PID_TEST_0 test_0.log
|
||||
|
||||
@@ -40,7 +40,9 @@ function cpu_tests() {
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
|
||||
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py"
|
||||
|
||||
# skip tests requiring model downloads if HF_TOKEN is not set
|
||||
# due to rate-limits
|
||||
@@ -97,3 +99,4 @@ function cpu_tests() {
|
||||
# All of CPU tests are expected to be finished less than 40 mins.
|
||||
export -f cpu_tests
|
||||
timeout 2h bash -c cpu_tests
|
||||
|
||||
|
||||
@@ -113,8 +113,8 @@ $PYTHON .buildkite/scripts/generate-nightly-index.py \
|
||||
echo "Uploading indices to $S3_COMMIT_PREFIX"
|
||||
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
|
||||
|
||||
# Update rocm/nightly/ if on main branch and not a PR
|
||||
if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]] || [[ "$NIGHTLY" == "1" ]]; then
|
||||
# Only scheduled nightly builds should update the moving nightly index.
|
||||
if [[ "${NIGHTLY:-0}" == "1" ]]; then
|
||||
echo "Updating rocm/nightly/ index..."
|
||||
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/rocm/nightly/"
|
||||
fi
|
||||
@@ -147,7 +147,7 @@ echo ""
|
||||
echo "Install command (by commit):"
|
||||
echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/$ROCM_SUBPATH/"
|
||||
echo ""
|
||||
if [[ "$BUILDKITE_BRANCH" == "main" ]] || [[ "$NIGHTLY" == "1" ]]; then
|
||||
if [[ "${NIGHTLY:-0}" == "1" ]]; then
|
||||
echo "Install command (nightly):"
|
||||
echo " pip install vllm --extra-index-url https://${BUCKET}.s3.amazonaws.com/rocm/nightly/"
|
||||
fi
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: Basic Correctness
|
||||
key: basic-correctness
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: Benchmarks CLI Test
|
||||
key: benchmarks-cli-test
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 45
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -18,6 +18,7 @@ steps:
|
||||
- pytest -v -s cuda/test_platform_no_cuda_init.py
|
||||
|
||||
- label: Cudagraph
|
||||
device: h200_35gb
|
||||
key: cudagraph
|
||||
timeout_in_minutes: 30
|
||||
source_file_dependencies:
|
||||
@@ -25,7 +26,10 @@ steps:
|
||||
- vllm/v1/cudagraph_dispatcher.py
|
||||
- vllm/config/compilation.py
|
||||
- vllm/compilation
|
||||
- vllm/v1/worker/encoder_cudagraph.py
|
||||
- vllm/v1/worker/encoder_cudagraph_defs.py
|
||||
commands:
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
|
||||
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
|
||||
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
|
||||
- pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
|
||||
|
||||
@@ -51,7 +51,7 @@ steps:
|
||||
|
||||
- label: e2e Scheduling (1 GPU)
|
||||
key: e2e-scheduling-1-gpu
|
||||
timeout_in_minutes: 35
|
||||
timeout_in_minutes: 53
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/v1/
|
||||
|
||||
@@ -3,6 +3,7 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Entrypoints Unit Tests
|
||||
device: h200_35gb
|
||||
key: entrypoints-unit-tests
|
||||
timeout_in_minutes: 25
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -15,6 +16,7 @@ steps:
|
||||
- pytest -v -s entrypoints/weight_transfer
|
||||
|
||||
- label: Entrypoints Integration (LLM)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-llm
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -37,7 +39,7 @@ steps:
|
||||
- label: Entrypoints Integration (API Server)
|
||||
key: entrypoints-integration-api-server
|
||||
device: h200_35gb
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -55,8 +57,9 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 1)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-1
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -73,8 +76,9 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server OpenAI - Part 2)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-openai-part-2
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 83
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -92,6 +96,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (API Server Generate)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-api-server-generate
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -114,6 +119,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Entrypoints Integration (Responses API)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-responses-api
|
||||
timeout_in_minutes: 50
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
@@ -148,8 +154,9 @@ steps:
|
||||
- pytest -v -s entrypoints/multimodal
|
||||
|
||||
- label: Entrypoints Integration (Pooling)
|
||||
device: h200_35gb
|
||||
key: entrypoints-integration-pooling
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
|
||||
@@ -15,6 +15,7 @@ steps:
|
||||
- pytest -v -s tests/kernels/ir
|
||||
|
||||
- label: Kernels Core Operation Test
|
||||
device: h200_35gb
|
||||
key: kernels-core-operation-test
|
||||
timeout_in_minutes: 120
|
||||
source_file_dependencies:
|
||||
@@ -163,6 +164,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Kernels Mamba Test
|
||||
device: h200_35gb
|
||||
key: kernels-mamba-test
|
||||
timeout_in_minutes: 40
|
||||
source_file_dependencies:
|
||||
@@ -235,6 +237,11 @@ steps:
|
||||
- vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py
|
||||
- vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py
|
||||
- vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py
|
||||
- vllm/cute_utils/
|
||||
- vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/
|
||||
- vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py
|
||||
- tests/kernels/mamba/test_gdn_prefill_cutedsl.py
|
||||
- tests/kernels/test_bf16x3_router_gemm_cutedsl.py
|
||||
- tests/kernels/test_ll_bf16_gemm.py
|
||||
- tests/kernels/test_top_k_per_row.py
|
||||
commands:
|
||||
@@ -264,6 +271,8 @@ steps:
|
||||
- pytest -v -s tests/kernels/moe/test_flashinfer_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_trtllm_nvfp4_moe.py
|
||||
- pytest -v -s tests/kernels/moe/test_cutedsl_moe.py
|
||||
- pytest -v -s tests/kernels/mamba/test_gdn_prefill_cutedsl.py
|
||||
- pytest -v -s tests/kernels/test_bf16x3_router_gemm_cutedsl.py
|
||||
- pytest -v -s tests/kernels/test_ll_bf16_gemm.py
|
||||
# e2e
|
||||
- pytest -v -s tests/models/quantization/test_nvfp4.py
|
||||
|
||||
@@ -78,6 +78,28 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt
|
||||
|
||||
- label: LM Eval PCP (4xB200)
|
||||
key: lm-eval-pcp-4xb200
|
||||
timeout_in_minutes: 360
|
||||
device: b200-k8s
|
||||
num_devices: 4
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
|
||||
- tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
|
||||
- tests/evals/gsm8k/configs/models-pcp.txt
|
||||
- vllm/model_executor/layers/quantization
|
||||
- vllm/config/parallel.py
|
||||
- vllm/distributed/parallel_state.py
|
||||
- vllm/model_executor/layers/attention/mla_attention.py
|
||||
- vllm/model_executor/layers/attention/pcp.py
|
||||
- vllm/v1/worker/gpu/model_runner.py
|
||||
- vllm/v1/worker/gpu/pcp_manager.py
|
||||
autorun_on_main: true
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt
|
||||
|
||||
- label: LM Eval Large Models EP (2xB200)
|
||||
key: lm-eval-large-models-ep-2xb200
|
||||
timeout_in_minutes: 60
|
||||
|
||||
@@ -31,7 +31,7 @@ steps:
|
||||
|
||||
- label: V1 Sample + Logits
|
||||
key: v1-sample-logits
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 83
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/config/
|
||||
@@ -64,8 +64,9 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: V1 Core + KV + Metrics
|
||||
device: h200_35gb
|
||||
key: v1-core-kv-metrics
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 80
|
||||
source_file_dependencies:
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
@@ -89,6 +90,7 @@ steps:
|
||||
- tests/v1/kv_offload
|
||||
- tests/v1/simple_kv_offload
|
||||
- tests/v1/worker
|
||||
- tests/v1/streaming_input
|
||||
- tests/v1/kv_connector/unit
|
||||
- tests/v1/ec_connector/unit
|
||||
- tests/v1/metrics
|
||||
@@ -102,6 +104,7 @@ steps:
|
||||
- pytest -v -s v1/kv_offload
|
||||
- pytest -v -s v1/simple_kv_offload
|
||||
- pytest -v -s v1/worker
|
||||
- pytest -v -s v1/streaming_input
|
||||
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
|
||||
- pytest -v -s -m 'not cpu_test' v1/metrics
|
||||
@@ -142,6 +145,7 @@ steps:
|
||||
- pytest -v -s -m 'cpu_test' v1/core
|
||||
- pytest -v -s v1/structured_output
|
||||
- pytest -v -s v1/test_serial_utils.py
|
||||
- pytest -v -s v1/cudagraph/test_cudagraph_manager.py
|
||||
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
|
||||
- pytest -v -s -m 'cpu_test' v1/metrics
|
||||
|
||||
|
||||
@@ -3,13 +3,16 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Model Executor
|
||||
device: h200_35gb
|
||||
key: model-executor
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 60
|
||||
source_file_dependencies:
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/config/model.py
|
||||
- vllm/model_executor
|
||||
- vllm/model_executor/warmup
|
||||
- tests/model_executor
|
||||
- tests/model_executor/test_jit_warmup.py
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
commands:
|
||||
- apt-get update && apt-get install -y curl libsodium23
|
||||
@@ -33,7 +36,9 @@ steps:
|
||||
- vllm/engine/arg_utils.py
|
||||
- vllm/config/model.py
|
||||
- vllm/model_executor
|
||||
- vllm/model_executor/warmup
|
||||
- tests/model_executor
|
||||
- tests/model_executor/test_jit_warmup.py
|
||||
- tests/entrypoints/openai/completion/test_tensorizer_entrypoint.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
@@ -46,6 +46,19 @@ steps:
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
|
||||
- label: Inkling Unit Tests (B200)
|
||||
key: inkling-unit-tests-b200
|
||||
timeout_in_minutes: 40
|
||||
device: b200-k8s
|
||||
source_file_dependencies:
|
||||
- vllm/models/inkling/
|
||||
- vllm/cute_utils/
|
||||
- cmake/external_projects/tml_fa4.cmake
|
||||
- tests/models/inkling/
|
||||
commands:
|
||||
# FA4 kernel tests require SM100; the suite skips them elsewhere.
|
||||
- pytest -v -s models/inkling
|
||||
|
||||
- label: Basic Models Test (Other CPU) # 5min
|
||||
key: basic-models-test-other-cpu
|
||||
depends_on:
|
||||
|
||||
@@ -21,6 +21,7 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: Language Models Tests (Extra Standard) %N
|
||||
device: h200_35gb
|
||||
key: language-models-tests-extra-standard
|
||||
timeout_in_minutes: 40
|
||||
source_file_dependencies:
|
||||
@@ -51,8 +52,8 @@ steps:
|
||||
- tests/models/language/pooling/test_classification.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/platforms/rocm.py
|
||||
|
||||
- label: Language Models Tests (Hybrid) %N
|
||||
device: h200_35gb
|
||||
key: language-models-tests-hybrid
|
||||
timeout_in_minutes: 65
|
||||
source_file_dependencies:
|
||||
@@ -63,8 +64,8 @@ steps:
|
||||
# Note: also needed to run plamo2 model in vLLM
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
# Shard hybrid language model tests
|
||||
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
# Shard the hybrid language model tests that are numerically stable on Hopper.
|
||||
- pytest -v -s models/language/generation -m hybrid_model -k 'not granite-4.0-tiny-preview' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
parallelism: 2
|
||||
mirror:
|
||||
amd:
|
||||
@@ -77,6 +78,20 @@ steps:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m hybrid_model --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB
|
||||
|
||||
# Granite 4 hybrid generation is sensitive to hardware-specific Triton SSD
|
||||
# autotuning (https://github.com/vllm-project/vllm/issues/25194). Keep this one
|
||||
# correctness test on L4 until its H200 output matches the Transformers reference.
|
||||
- label: Language Models Tests (Granite L4 Compatibility)
|
||||
key: language-models-tests-granite-l4-compatibility
|
||||
timeout_in_minutes: 65
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/language/generation
|
||||
commands:
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/state-spaces/mamba@v2.3.0'
|
||||
- uv pip install --system --no-build-isolation 'git+https://github.com/Dao-AILab/causal-conv1d@v1.6.0'
|
||||
- pytest -v -s models/language/generation -m hybrid_model -k 'granite-4.0-tiny-preview'
|
||||
|
||||
- label: Language Models Test (Extended Generation) # 80min
|
||||
device: h200_35gb
|
||||
key: language-models-test-extended-generation
|
||||
@@ -122,7 +137,7 @@ steps:
|
||||
|
||||
- label: Language Models Test (MTEB)
|
||||
key: language-models-test-mteb
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
device: h200_18gb
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -4,7 +4,7 @@ depends_on:
|
||||
steps:
|
||||
- label: "Multi-Modal Models (Standard) 1: qwen2"
|
||||
key: multi-modal-models-standard-1-qwen2
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 68
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -20,7 +20,7 @@ steps:
|
||||
|
||||
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
|
||||
key: multi-modal-models-standard-2-qwen3-gemma
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -54,7 +54,7 @@ steps:
|
||||
- label: "Multi-Modal Models (Standard) 4: other + whisper"
|
||||
device: h200_35gb
|
||||
key: multi-modal-models-standard-4-other-whisper
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 75
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
@@ -85,7 +85,7 @@ steps:
|
||||
|
||||
- label: Multi-Modal Processor # 44min
|
||||
key: multi-modal-processor
|
||||
timeout_in_minutes: 65
|
||||
timeout_in_minutes: 98
|
||||
device: h200_18gb
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
@@ -119,6 +119,7 @@ steps:
|
||||
- vllm/model_executor/model_loader/
|
||||
|
||||
- label: Multi-Modal Models (Extended Generation 1)
|
||||
device: h200_35gb
|
||||
key: multi-modal-models-extended-generation-1
|
||||
optional: true
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -5,7 +5,7 @@ steps:
|
||||
- label: PyTorch Compilation Unit Tests
|
||||
device: h200_35gb
|
||||
key: pytorch-compilation-unit-tests
|
||||
timeout_in_minutes: 90
|
||||
timeout_in_minutes: 150
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
@@ -116,8 +116,9 @@ steps:
|
||||
- image-build-amd
|
||||
|
||||
- label: PyTorch Fullgraph Smoke Test
|
||||
device: h200_35gb
|
||||
key: pytorch-fullgraph-smoke-test
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 90
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
@@ -149,7 +150,42 @@ steps:
|
||||
# as it is a heavy test that is covered in other steps.
|
||||
# Use `find` to launch multiple instances of pytest so that
|
||||
# they do not suffer from https://github.com/vllm-project/vllm/issues/28965
|
||||
- "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
|
||||
- "find compile/fullgraph/ -name 'test_*.py' -not -name 'test_full_cudagraph.py' -not -name 'test_full_graph.py' -print0 | xargs -0 -n1 -I{} pytest -s -v '{}'"
|
||||
|
||||
# Hopper-only DeepSeek-V2-Lite cases in this file require two 29.3-GiB model
|
||||
# instances and cannot fit a 35GB MIG slice. L4 retains the original coverage:
|
||||
# those SM90 cases skip while the architecture-compatible cases still run.
|
||||
- label: PyTorch Fullgraph CUDAGraph (L4 Compatibility)
|
||||
key: pytorch-fullgraph-cudagraph-l4-compatibility
|
||||
timeout_in_minutes: 60
|
||||
source_file_dependencies:
|
||||
- vllm/__init__.py
|
||||
- vllm/_aiter_ops.py
|
||||
- vllm/_custom_ops.py
|
||||
- vllm/compilation/
|
||||
- vllm/config/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/env_override.py
|
||||
- vllm/envs.py
|
||||
- vllm/forward_context.py
|
||||
- vllm/inputs/
|
||||
- vllm/ir/
|
||||
- vllm/kernels/
|
||||
- vllm/logger.py
|
||||
- vllm/model_executor/
|
||||
- vllm/multimodal/
|
||||
- vllm/platforms/
|
||||
- vllm/plugins/
|
||||
- vllm/sampling_params.py
|
||||
- vllm/sequence.py
|
||||
- vllm/transformers_utils/
|
||||
- vllm/triton_utils/
|
||||
- vllm/utils/
|
||||
- vllm/v1/
|
||||
- tests/compile
|
||||
commands:
|
||||
- pytest -s -v compile/fullgraph/test_full_cudagraph.py
|
||||
|
||||
- label: PyTorch Fullgraph
|
||||
key: pytorch-fullgraph
|
||||
|
||||
@@ -3,8 +3,11 @@ depends_on:
|
||||
- image-build
|
||||
steps:
|
||||
- label: Quantization
|
||||
device: h200_35gb
|
||||
key: quantization
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 75
|
||||
env:
|
||||
VLLM_USE_V2_MODEL_RUNNER: "0"
|
||||
source_file_dependencies:
|
||||
- csrc/
|
||||
- vllm/model_executor/layers/quantization
|
||||
@@ -19,9 +22,13 @@ steps:
|
||||
# TODO(jerryzh168): resolve the above comment
|
||||
- uv pip install --system torchao==0.17.0 --index-url https://download.pytorch.org/whl/cu130
|
||||
- uv pip install --system conch-triton-kernels
|
||||
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
|
||||
# The SM90-only checkpoint currently contains a removed weight_chan_scale
|
||||
# parameter. It was not exercised by the previous L4 job.
|
||||
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py -k 'not test_compressed_tensors_w4a8_fp8' --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
parallelism: 8
|
||||
|
||||
- label: Quantized Fusions
|
||||
device: h200_35gb
|
||||
key: quantized-fusions
|
||||
timeout_in_minutes: 20
|
||||
source_file_dependencies:
|
||||
@@ -52,10 +59,14 @@ steps:
|
||||
- pytest -s -v tests/quantization/test_blackwell_moe.py
|
||||
|
||||
- label: Quantized Models Test
|
||||
device: h200_35gb
|
||||
key: quantized-models-test
|
||||
timeout_in_minutes: 50
|
||||
timeout_in_minutes: 65
|
||||
env:
|
||||
VLLM_USE_V2_MODEL_RUNNER: "0"
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization
|
||||
- tests/models/quantization
|
||||
commands:
|
||||
- pytest -v -s models/quantization
|
||||
- pytest -v -s models/quantization --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
parallelism: 3
|
||||
|
||||
@@ -81,6 +81,7 @@ steps:
|
||||
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
|
||||
- label: Rust Frontend Tool Use
|
||||
device: h200_35gb
|
||||
timeout_in_minutes: 25
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
|
||||
@@ -170,3 +170,19 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
|
||||
|
||||
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
|
||||
key: spec-decode-deepseek-mtp-parallel-load-b200
|
||||
timeout_in_minutes: 30
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 2
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/llm_base_proposer.py
|
||||
- vllm/v1/spec_decode/eagle.py
|
||||
- vllm/v1/worker/gpu/spec_decode/eagle/
|
||||
- vllm/model_executor/models/deepseek_mtp.py
|
||||
- vllm/model_executor/models/deepseek_v2.py
|
||||
- tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
|
||||
|
||||
@@ -47,6 +47,7 @@
|
||||
|
||||
# Rust Frontend
|
||||
/rust/ @BugenZhao @njhill
|
||||
/rust/src/bench @esmeetu
|
||||
/build_rust.sh @BugenZhao @njhill
|
||||
/rust-toolchain.toml @BugenZhao @njhill
|
||||
/.buildkite/test_areas/rust* @BugenZhao @njhill
|
||||
|
||||
@@ -181,6 +181,18 @@ pull_request_rules:
|
||||
add:
|
||||
- performance
|
||||
|
||||
- name: label-quantization
|
||||
description: Automatically apply quantization label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^vllm/model_executor/layers/quantization/
|
||||
- title~=(?i)quant
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- quantization
|
||||
|
||||
- name: label-qwen
|
||||
description: Automatically apply qwen label
|
||||
conditions:
|
||||
|
||||
@@ -130,6 +130,47 @@ jobs:
|
||||
},
|
||||
],
|
||||
},
|
||||
quantization: {
|
||||
keywords: [
|
||||
{
|
||||
term: "quantization",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "quantized",
|
||||
searchIn: "both"
|
||||
},
|
||||
],
|
||||
},
|
||||
"intel-gpu": {
|
||||
// Keyword search - matches whole words only (with word boundaries)
|
||||
keywords: [
|
||||
{
|
||||
term: "B50",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "B60",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "B70",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "intel gpu",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "Arc GPU",
|
||||
searchIn: "both"
|
||||
},
|
||||
{
|
||||
term: "BMG",
|
||||
searchIn: "both"
|
||||
},
|
||||
],
|
||||
},
|
||||
// Add more label configurations here as needed
|
||||
// example: {
|
||||
// keywords: [...],
|
||||
@@ -491,4 +532,4 @@ jobs:
|
||||
issue_number: context.issue.number,
|
||||
body: message,
|
||||
});
|
||||
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
|
||||
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
|
||||
|
||||
@@ -30,7 +30,7 @@ repos:
|
||||
- id: markdownlint-cli2
|
||||
language_version: lts
|
||||
args: [--fix]
|
||||
exclude: ^CLAUDE\.md$
|
||||
exclude: (^|/)CLAUDE\.md$
|
||||
- repo: https://github.com/rhysd/actionlint
|
||||
rev: v1.7.7
|
||||
hooks:
|
||||
|
||||
+2
-2
@@ -68,8 +68,8 @@ endif()
|
||||
# requirements.txt files and should be kept consistent. The ROCm torch
|
||||
# versions are derived from docker/Dockerfile.rocm
|
||||
#
|
||||
set(TORCH_SUPPORTED_VERSION_CUDA "2.11.0")
|
||||
set(TORCH_SUPPORTED_VERSION_ROCM "2.11.0")
|
||||
set(TORCH_SUPPORTED_VERSION_CUDA "2.13.0")
|
||||
set(TORCH_SUPPORTED_VERSION_ROCM "2.13.0")
|
||||
# TORCH_NIGHTLY=1 builds run against unpinned nightly wheels, so the supported-
|
||||
# version check would always warn. Only treat it as a nightly build when the
|
||||
# value is exactly "1" (the bootstrap exports TORCH_NIGHTLY=0 by default, which
|
||||
|
||||
@@ -48,7 +48,7 @@ vLLM is flexible and easy to use with:
|
||||
- Tool calling and reasoning parsers
|
||||
- OpenAI-compatible API server, plus Anthropic Messages API and gRPC support
|
||||
- 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.
|
||||
- Support for NVIDIA GPUs, AMD GPUs, Intel 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:
|
||||
|
||||
|
||||
@@ -69,12 +69,11 @@ def make_inputs(total_tokens, num_reqs, block_size):
|
||||
# Output workspace
|
||||
dst = torch.zeros(total_tokens, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
|
||||
|
||||
seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device="cuda")
|
||||
workspace_starts_t = torch.tensor(
|
||||
workspace_starts, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
|
||||
return cache, dst, block_table, seq_lens_t, workspace_starts_t
|
||||
return cache, dst, block_table, workspace_starts_t
|
||||
|
||||
|
||||
def bench_scenario(label, num_reqs, total_tokens_list, save_path):
|
||||
@@ -94,7 +93,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):
|
||||
)
|
||||
)
|
||||
def bench_fn(total_tokens, provider, num_reqs):
|
||||
cache, dst, block_table, seq_lens_t, ws_starts = make_inputs(
|
||||
cache, dst, block_table, ws_starts = make_inputs(
|
||||
total_tokens, num_reqs, BLOCK_SIZE
|
||||
)
|
||||
|
||||
@@ -102,7 +101,7 @@ def bench_scenario(label, num_reqs, total_tokens_list, save_path):
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: ops.cp_gather_and_upconvert_fp8_kv_cache(
|
||||
cache, dst, block_table, seq_lens_t, ws_starts, num_reqs
|
||||
cache, dst, block_table, ws_starts, num_reqs
|
||||
),
|
||||
quantiles=quantiles,
|
||||
rep=500,
|
||||
|
||||
@@ -430,6 +430,7 @@ set(VLLM_EXT_SRC
|
||||
"csrc/cpu/layernorm.cpp"
|
||||
"csrc/cpu/mla_decode.cpp"
|
||||
"csrc/cpu/pos_encoding.cpp"
|
||||
"csrc/cpu/mamba_cpu.cpp"
|
||||
"csrc/moe/dynamic_4bit_int_moe_cpu.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp")
|
||||
@@ -489,6 +490,7 @@ if (ENABLE_X86_ISA)
|
||||
"csrc/cpu/spec_decode_utils.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/dnnl_kernels.cpp"
|
||||
"csrc/cpu/mamba_cpu.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp"
|
||||
# TODO: Remove these files
|
||||
"csrc/cpu/activation.cpp"
|
||||
@@ -502,6 +504,7 @@ if (ENABLE_X86_ISA)
|
||||
"csrc/cpu/utils.cpp"
|
||||
"csrc/cpu/spec_decode_utils.cpp"
|
||||
"csrc/cpu/cpu_attn.cpp"
|
||||
"csrc/cpu/mamba_cpu.cpp"
|
||||
"csrc/cpu/dnnl_kernels.cpp"
|
||||
"csrc/cpu/torch_bindings.cpp"
|
||||
# TODO: Remove these files
|
||||
|
||||
@@ -17,7 +17,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
fmha_sm100
|
||||
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
|
||||
GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
|
||||
GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
|
||||
@@ -22,7 +22,7 @@ if(QUTLASS_SRC_DIR)
|
||||
set(qutlass_BINARY_DIR "${CMAKE_BINARY_DIR}/qutlass-binary-dir-unused")
|
||||
else()
|
||||
set(_QUTLASS_UPSTREAM_REPO "https://github.com/IST-DASLab/qutlass.git")
|
||||
set(_QUTLASS_UPSTREAM_TAG "830d2c4537c7396e14a02a46fbddd18b5d107c65")
|
||||
set(_QUTLASS_UPSTREAM_TAG "e74319e3405ce6d71965732880f5dc1f52371f64")
|
||||
|
||||
set(_qutlass_fc_root "${FETCHCONTENT_BASE_DIR}")
|
||||
if(NOT _qutlass_fc_root)
|
||||
@@ -125,8 +125,6 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
|
||||
CUDA_ARCHS "${QUTLASS_ARCHS}"
|
||||
)
|
||||
|
||||
# QuTLASS uses legacy ATen headers and cannot be built with TORCH_TARGET_VERSION.
|
||||
# Keep it as its own extension (registers torch.ops._qutlass_C).
|
||||
define_extension_target(
|
||||
_qutlass_C
|
||||
DESTINATION vllm
|
||||
@@ -139,9 +137,11 @@ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.8 AND QUTLASS_ARCHS)
|
||||
WITH_SOABI)
|
||||
|
||||
target_compile_definitions(_qutlass_C PRIVATE
|
||||
QUTLASS_DISABLE_PYBIND=1
|
||||
QUTLASS_MINIMAL_BUILD=1
|
||||
TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)
|
||||
CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1
|
||||
TORCH_TARGET_VERSION=0x020B000000000000ULL
|
||||
USE_CUDA)
|
||||
|
||||
set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
|
||||
|
||||
@@ -14,7 +14,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
tml_fa4
|
||||
GIT_REPOSITORY https://github.com/vllm-project/tml-fa4.git
|
||||
GIT_TAG 13374f0c855acc1add1bf30444bd67aebbc24a8e
|
||||
GIT_TAG b206834606ed5b5f21f8eed6b0683f528ea9cf7d
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND "")
|
||||
|
||||
@@ -39,7 +39,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
vllm-flash-attn
|
||||
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
|
||||
GIT_TAG caaa4eb59845388a20b1f435ecaafb4bd9517ad8
|
||||
GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32
|
||||
GIT_PROGRESS TRUE
|
||||
# Don't share the vllm-flash-attn build between build types
|
||||
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
|
||||
|
||||
+1
-2
@@ -67,9 +67,8 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]
|
||||
torch::Tensor const& dst, // [TOT_TOKENS, 576]
|
||||
torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
|
||||
torch::Tensor const& seq_lens, // [BATCH]
|
||||
torch::Tensor const& workspace_starts, // [BATCH]
|
||||
int64_t batch_size);
|
||||
int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
|
||||
|
||||
// Indexer K quantization and cache function
|
||||
void indexer_k_quant_and_cache(
|
||||
|
||||
@@ -102,7 +102,9 @@ class TileGemm82 {
|
||||
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);
|
||||
auto fp32_b_regs = load_b_pair_vec(curr_b);
|
||||
auto fp32_b_0_reg = fp32_b_regs.first;
|
||||
auto fp32_b_1_reg = fp32_b_regs.second;
|
||||
|
||||
float* __restrict__ curr_m_a = curr_a;
|
||||
vec_op::unroll_loop<int32_t, M>([&](int32_t i) {
|
||||
|
||||
@@ -336,13 +336,14 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
reg.val[1] = fp16_to_fp32_bits(raw_lo);
|
||||
}
|
||||
float reduce_sum() const {
|
||||
AliasReg ar;
|
||||
ar.reg = reg;
|
||||
float result = 0;
|
||||
unroll_loop<int, VEC_ELEM_NUM>(
|
||||
[&result, &ar](int i) { result += ar.values[i]; });
|
||||
|
||||
return result;
|
||||
// VSX horizontal reduction: 3 vector ops instead of 8 scalar adds.
|
||||
// Step 1: pairwise sum of the two 4-wide halves
|
||||
__vector float s = vec_add(reg.val[0], reg.val[1]);
|
||||
// Step 2: rotate by 8 bytes (2 floats) and add
|
||||
s = vec_add(s, vec_sld(s, s, 8));
|
||||
// Step 3: rotate by 4 bytes (1 float) and add => all lanes hold total
|
||||
s = vec_add(s, vec_sld(s, s, 4));
|
||||
return vec_extract(s, 0);
|
||||
}
|
||||
FP32Vec8 exp() const {
|
||||
f32x4x2_t out;
|
||||
|
||||
@@ -0,0 +1,285 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// CPU at::Tensor wrappers for Mamba decode-step kernels defined in
|
||||
// mamba_kernels.hpp.
|
||||
|
||||
#include "cpu/mamba_kernels.hpp"
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <torch/library.h>
|
||||
#include <c10/util/Optional.h>
|
||||
|
||||
#include "cpu_types.hpp"
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// causal_conv1d_update
|
||||
// ---------------------------------------------------------------------------
|
||||
at::Tensor causal_conv1d_update_cpu_impl(
|
||||
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
|
||||
const c10::optional<at::Tensor>& bias,
|
||||
const c10::optional<std::string>& activation,
|
||||
const c10::optional<at::Tensor>& conv_state_indices,
|
||||
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id) {
|
||||
bool do_silu = false;
|
||||
if (activation.has_value()) {
|
||||
const std::string& act = activation.value();
|
||||
do_silu = (act == "silu" || act == "swish");
|
||||
}
|
||||
|
||||
at::ScalarType dtype = x.scalar_type();
|
||||
|
||||
// Input x: contiguous in native dtype.
|
||||
at::Tensor x_c = x.is_contiguous() ? x : x.contiguous();
|
||||
|
||||
// conv_state: NEVER copy the full paged tensor just for layout reasons.
|
||||
// If the dtype matches we work directly on conv_state (contiguous or not)
|
||||
// by extracting strides and passing them to the kernel.
|
||||
// Only a dtype-conversion copy is made when types differ (rare for BF16).
|
||||
bool state_type_ok = (conv_state.scalar_type() == dtype);
|
||||
at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype);
|
||||
// state_c and conv_state may be non-contiguous — that is intentional.
|
||||
|
||||
// Weight: coerce to same dtype if needed (should match in practice)
|
||||
at::Tensor w_c =
|
||||
(weight.scalar_type() != dtype)
|
||||
? weight.to(dtype).contiguous()
|
||||
: (weight.is_contiguous() ? weight : weight.contiguous());
|
||||
|
||||
// Bias stays float32 (small scalar, used only for fp32 accumulation)
|
||||
at::Tensor bias_f32;
|
||||
if (bias.has_value() && bias.value().defined())
|
||||
bias_f32 = bias.value().to(at::kFloat).contiguous();
|
||||
|
||||
int64_t batch = x_c.size(0);
|
||||
int64_t dim = x_c.size(1);
|
||||
int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1;
|
||||
int64_t width = w_c.size(1);
|
||||
int64_t state_len = state_c.size(2);
|
||||
|
||||
// Extract strides — works for contiguous AND non-contiguous (transposed)
|
||||
// state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in
|
||||
// contiguous) stride(1): between conv channels (dim stride) stride(2):
|
||||
// between state elements (=1 when contiguous, =dim when transposed)
|
||||
int64_t stride_s_slot = state_c.stride(0);
|
||||
int64_t stride_s_dim = state_c.stride(1);
|
||||
int64_t stride_s_state = state_c.stride(2);
|
||||
|
||||
at::Tensor out = x_c.clone(); // native dtype, no float32 alloc
|
||||
|
||||
const int32_t* cache_idx_ptr = nullptr;
|
||||
at::Tensor cache_idx_int;
|
||||
if (conv_state_indices.has_value()) {
|
||||
cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous();
|
||||
cache_idx_ptr = cache_idx_int.data_ptr<int32_t>();
|
||||
}
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] {
|
||||
mamba_cpu::causal_conv1d_update_kernel<scalar_t>(
|
||||
x_c.data_ptr<scalar_t>(), state_c.data_ptr<scalar_t>(), stride_s_slot,
|
||||
stride_s_dim, stride_s_state, w_c.data_ptr<scalar_t>(),
|
||||
bias_f32.defined() ? bias_f32.data_ptr<float>() : nullptr,
|
||||
out.data_ptr<scalar_t>(), cache_idx_ptr,
|
||||
static_cast<int32_t>(pad_slot_id), batch, dim, seqlen, width, state_len,
|
||||
do_silu);
|
||||
});
|
||||
|
||||
// Write back only when a type-conversion copy was made.
|
||||
// Layout-only non-contiguity is handled via strides above — no copy needed.
|
||||
if (!state_type_ok) conv_state.copy_(state_c);
|
||||
|
||||
return out;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// selective_state_update
|
||||
// ---------------------------------------------------------------------------
|
||||
void selective_state_update_cpu_impl(
|
||||
at::Tensor& state, // (nstates, nheads, dim, dstate)
|
||||
const at::Tensor& x, // (N, nheads, dim)
|
||||
const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B,
|
||||
const at::Tensor& C, const c10::optional<at::Tensor>& D,
|
||||
const c10::optional<at::Tensor>& z,
|
||||
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
|
||||
const c10::optional<at::Tensor>& state_batch_indices,
|
||||
const c10::optional<at::Tensor>& dst_state_batch_indices,
|
||||
int64_t null_block_id, at::Tensor& out,
|
||||
const c10::optional<at::Tensor>& num_accepted_tokens,
|
||||
const c10::optional<at::Tensor>& cu_seqlens) {
|
||||
at::ScalarType state_type = state.scalar_type();
|
||||
at::ScalarType input_type = x.scalar_type();
|
||||
|
||||
// x, B, C must be contiguous and match input_type
|
||||
auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
at::Tensor x_in = ensure_input(x);
|
||||
at::Tensor B_in = ensure_input(B);
|
||||
at::Tensor C_in = ensure_input(C);
|
||||
at::Tensor z_in;
|
||||
if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value());
|
||||
|
||||
// A, D, dt_bias are float32 model parameters that arrive here as expanded
|
||||
// tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0).
|
||||
// We need just the scalar value per head as a (nheads,) 1-D array so that
|
||||
// A_ptr[h] in the kernel correctly reads head h's value.
|
||||
//
|
||||
// Strategy: peel trailing expanded (stride=0) dims via .select(), which is
|
||||
// a zero-copy view. For A: (nheads, head_dim, dstate) strides (1,0,0)
|
||||
// → .select(2,0) → (nheads, head_dim) strides (1,0)
|
||||
// → .select(1,0) → (nheads,) stride (1,) ← contiguous, free.
|
||||
// No allocation, no type conversion (A is already float32).
|
||||
auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = t;
|
||||
// Peel trailing dimensions that are broadcast (stride=0 or size=1)
|
||||
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
|
||||
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
|
||||
at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32
|
||||
at::Tensor D_f32, dt_bias_f32;
|
||||
if (D.has_value() && D.value().defined())
|
||||
D_f32 = to_per_head_1d_f32(D.value());
|
||||
if (dt_bias.has_value() && dt_bias.value().defined())
|
||||
dt_bias_f32 = to_per_head_1d_f32(dt_bias.value());
|
||||
|
||||
// dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE
|
||||
// the type conversion so we convert head_dim x fewer elements.
|
||||
at::Tensor dt_f32;
|
||||
{
|
||||
// If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2,
|
||||
// take a zero-copy view of index 0 along that dim first.
|
||||
at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt; // (N, nheads)
|
||||
at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2;
|
||||
dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous();
|
||||
}
|
||||
|
||||
int64_t nheads = state.size(1);
|
||||
int64_t dim = state.size(2);
|
||||
int64_t dstate = state.size(3);
|
||||
int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined())
|
||||
? cu_seqlens.value().size(0) - 1
|
||||
: x_in.size(0);
|
||||
int64_t ngroups = B_in.size(1);
|
||||
|
||||
// Strides
|
||||
int64_t stride_state_n = state.stride(0);
|
||||
int64_t stride_state_h = state.stride(1);
|
||||
int64_t stride_state_d = state.stride(2);
|
||||
int64_t stride_x_n = x_in.stride(0);
|
||||
int64_t stride_x_h = x_in.stride(1);
|
||||
int64_t stride_dt_n = dt_f32.stride(0); // dt is (N, nheads)
|
||||
int64_t stride_BC_n = B_in.stride(0);
|
||||
int64_t stride_BC_g = B_in.stride(1);
|
||||
int64_t stride_out_n = out.stride(0);
|
||||
int64_t stride_out_h = out.stride(1);
|
||||
|
||||
// Optional index pointers
|
||||
auto get_int32_ptr =
|
||||
[](const c10::optional<at::Tensor>& opt) -> const int32_t* {
|
||||
return (opt.has_value() && opt.value().defined())
|
||||
? opt.value().data_ptr<int32_t>()
|
||||
: nullptr;
|
||||
};
|
||||
const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices);
|
||||
const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices);
|
||||
const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens);
|
||||
const int32_t* csl_ptr = get_int32_ptr(cu_seqlens);
|
||||
|
||||
// Dispatch on (state_t, input_t, out_t): write directly into `out`
|
||||
// without any intermediate float32 buffer.
|
||||
VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] {
|
||||
using state_t = scalar_t;
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] {
|
||||
using input_t = scalar_t;
|
||||
VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] {
|
||||
using out_t = scalar_t;
|
||||
mamba_cpu::selective_state_update_kernel<state_t, input_t, out_t>(
|
||||
state.data_ptr<state_t>(), stride_state_n, stride_state_h,
|
||||
stride_state_d, x_in.data_ptr<input_t>(), stride_x_n, stride_x_h,
|
||||
dt_f32.data_ptr<float>(), stride_dt_n, A_f32.data_ptr<float>(),
|
||||
B_in.data_ptr<input_t>(), C_in.data_ptr<input_t>(), stride_BC_n,
|
||||
stride_BC_g, D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
|
||||
z_in.defined() ? z_in.data_ptr<input_t>() : nullptr,
|
||||
dt_bias_f32.defined() ? dt_bias_f32.data_ptr<float>() : nullptr,
|
||||
out.data_ptr<out_t>(), stride_out_n, stride_out_h, sbi_ptr,
|
||||
dsbi_ptr, static_cast<int32_t>(null_block_id), nat_ptr, csl_ptr, N,
|
||||
nheads, ngroups, dim, dstate, dt_softplus);
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mamba_chunk_scan_fwd_cpu
|
||||
// ---------------------------------------------------------------------------
|
||||
void mamba_chunk_scan_fwd_cpu_impl(
|
||||
at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller
|
||||
at::Tensor&
|
||||
final_states, // [batch, nheads, headdim, dstate] float32 contiguous
|
||||
const at::Tensor& x, // [seqlen, nheads, headdim]
|
||||
const at::Tensor&
|
||||
dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp)
|
||||
const at::Tensor& A, // [nheads] float32
|
||||
const at::Tensor& B, // [seqlen, ngroups, dstate]
|
||||
const at::Tensor& C, // [seqlen, ngroups, dstate]
|
||||
const c10::optional<at::Tensor>& D, // [nheads] float32 (optional)
|
||||
const c10::optional<at::Tensor>& z, // [seqlen, nheads, headdim] (optional)
|
||||
const at::Tensor& cu_seqlens // [batch+1] int32
|
||||
) {
|
||||
const at::ScalarType input_type = x.scalar_type();
|
||||
|
||||
auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t;
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
at::Tensor x_in = ensure_contig(x);
|
||||
at::Tensor B_in = ensure_contig(B);
|
||||
at::Tensor C_in = ensure_contig(C);
|
||||
at::Tensor z_in;
|
||||
if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value());
|
||||
|
||||
// A and D are float32 model parameters, potentially broadcast-expanded.
|
||||
// Strip trailing broadcast dims to get a contiguous (nheads,) array.
|
||||
auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor {
|
||||
at::Tensor r = t;
|
||||
while (r.dim() > 1) r = r.select(r.dim() - 1, 0);
|
||||
if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat);
|
||||
return r.is_contiguous() ? r : r.contiguous();
|
||||
};
|
||||
at::Tensor A_f32 = to_per_head_f32(A);
|
||||
at::Tensor D_f32;
|
||||
if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value());
|
||||
|
||||
// dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in
|
||||
// Python.
|
||||
at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous();
|
||||
if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat);
|
||||
|
||||
at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous();
|
||||
|
||||
const int64_t batch = final_states.size(0);
|
||||
const int64_t nheads = final_states.size(1);
|
||||
const int64_t headdim = final_states.size(2);
|
||||
const int64_t dstate = final_states.size(3);
|
||||
const int64_t ngroups = B_in.size(1);
|
||||
|
||||
TORCH_CHECK(final_states.is_contiguous(),
|
||||
"mamba_chunk_scan_fwd_cpu: final_states must be contiguous");
|
||||
TORCH_CHECK(out.is_contiguous(),
|
||||
"mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via "
|
||||
"raw data_ptr)");
|
||||
|
||||
VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] {
|
||||
mamba_cpu::mamba_chunk_scan_fwd_kernel<scalar_t>(
|
||||
final_states.data_ptr<float>(), x_in.data_ptr<scalar_t>(),
|
||||
dt_c.data_ptr<float>(), A_f32.data_ptr<float>(),
|
||||
B_in.data_ptr<scalar_t>(), C_in.data_ptr<scalar_t>(),
|
||||
D_f32.defined() ? D_f32.data_ptr<float>() : nullptr,
|
||||
z_in.defined() ? z_in.data_ptr<scalar_t>() : nullptr,
|
||||
out.data_ptr<scalar_t>(), cu_int.data_ptr<int32_t>(), batch, nheads,
|
||||
ngroups, headdim, dstate);
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,382 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
//
|
||||
// Fused CPU vector kernels for Mamba decode-step hotspots:
|
||||
// - causal_conv1d_update (depthwise 1-D conv state roll + compute)
|
||||
// - selective_state_update (SSM recurrence, single-step)
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "cpu_types.hpp"
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <cstdint>
|
||||
#include <algorithm>
|
||||
|
||||
namespace mamba_cpu {
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// causal_conv1d_update — templated for native BF16/FP32
|
||||
//
|
||||
// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor.
|
||||
// Explicit strides are passed so the kernel writes directly into the
|
||||
// correct memory locations without making a contiguous copy of the full
|
||||
// paged tensor (which was the source of the 34-41% direct_copy_kernel).
|
||||
//
|
||||
// stride_s_slot = state.stride(0) — between cache slots
|
||||
// stride_s_dim = state.stride(1) — between conv_dim channels
|
||||
// stride_s_state = state.stride(2) — between state elements
|
||||
//
|
||||
// When stride_s_state == 1 (contiguous), the memmove fast path is used.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename scalar_t>
|
||||
inline void causal_conv1d_update_kernel(
|
||||
const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr,
|
||||
int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state,
|
||||
const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr,
|
||||
scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs,
|
||||
int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen,
|
||||
int64_t width, int64_t state_len, bool do_silu) {
|
||||
#pragma omp parallel for
|
||||
for (int64_t b = 0; b < batch; ++b) {
|
||||
int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b;
|
||||
if (cache_idx == pad_slot_id) continue;
|
||||
|
||||
for (int64_t t = 0; t < seqlen; ++t) {
|
||||
const scalar_t* x_b = x_ptr + (b * dim * seqlen + t);
|
||||
scalar_t* out_b = out_ptr + (b * dim * seqlen + t);
|
||||
// Base of this slot in the (possibly non-contiguous) paged state
|
||||
scalar_t* s_base = state_ptr + cache_idx * stride_s_slot;
|
||||
|
||||
for (int64_t d = 0; d < dim; ++d) {
|
||||
float x_val = static_cast<float>(x_b[d * seqlen]);
|
||||
scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state
|
||||
const scalar_t* w = weight_ptr + d * width;
|
||||
|
||||
// Accumulate in float32 for precision
|
||||
float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f;
|
||||
for (int64_t k = 0; k < state_len; ++k) {
|
||||
acc += static_cast<float>(w[k]) *
|
||||
static_cast<float>(sd[k * stride_s_state]);
|
||||
}
|
||||
acc += static_cast<float>(w[state_len]) * x_val;
|
||||
|
||||
// Shift state left and append new input.
|
||||
// Use memmove when contiguous (stride==1); element loop otherwise.
|
||||
if (stride_s_state == 1) {
|
||||
if (state_len > 1)
|
||||
std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t));
|
||||
if (state_len > 0) sd[state_len - 1] = static_cast<scalar_t>(x_val);
|
||||
} else {
|
||||
for (int64_t k = 0; k < state_len - 1; ++k)
|
||||
sd[k * stride_s_state] = sd[(k + 1) * stride_s_state];
|
||||
if (state_len > 0)
|
||||
sd[(state_len - 1) * stride_s_state] = static_cast<scalar_t>(x_val);
|
||||
}
|
||||
|
||||
if (do_silu) {
|
||||
float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc))
|
||||
: std::exp(acc) / (1.0f + std::exp(acc));
|
||||
acc *= sigmoid;
|
||||
}
|
||||
out_b[d * seqlen] = static_cast<scalar_t>(acc);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// selective_state_update
|
||||
//
|
||||
// Template parameters:
|
||||
// state_t - dtype of ssm_state cache (typically BFloat16)
|
||||
// input_t - dtype of x, B, C (typically BFloat16)
|
||||
// out_t - dtype of output tensor (typically BFloat16)
|
||||
// Write directly — no float32 intermediate buffer needed.
|
||||
//
|
||||
// A, D, dt_bias are accepted as const float* (they are always float32
|
||||
// model parameters in Mamba2). This eliminates the per-call float32→BF16
|
||||
// conversion and the .contiguous() materialisation of the broadcast-expand.
|
||||
//
|
||||
// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the
|
||||
// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed.
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename state_t, typename input_t, typename out_t = float>
|
||||
inline void selective_state_update_kernel(
|
||||
state_t* __restrict__ state_ptr, int64_t stride_state_n,
|
||||
int64_t stride_state_h, int64_t stride_state_d,
|
||||
const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h,
|
||||
// dt: (N, nheads) — scalar per head, NOT expanded to head_dim
|
||||
const float* __restrict__ dt_ptr, int64_t stride_dt_n,
|
||||
// A: (nheads,) float32 — scalar per head
|
||||
const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr,
|
||||
const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g,
|
||||
// D: (nheads,) float32 — scalar per head (nullptr if not used)
|
||||
const float* __restrict__ D_ptr,
|
||||
// z: same shape as x (optional)
|
||||
const input_t* __restrict__ z_ptr,
|
||||
// dt_bias: (nheads,) float32 — scalar per head (nullptr if not used)
|
||||
const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr,
|
||||
int64_t stride_out_n, int64_t stride_out_h,
|
||||
const int32_t* __restrict__ state_batch_indices,
|
||||
const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id,
|
||||
const int32_t* __restrict__ num_accepted_tokens,
|
||||
const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads,
|
||||
int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) {
|
||||
using state_vec_t = vec_op::vec_t<state_t>;
|
||||
using input_vec_t = vec_op::vec_t<input_t>;
|
||||
constexpr int VEC_ELEM_NUM = 8;
|
||||
|
||||
int64_t nheads_per_group = nheads / ngroups;
|
||||
|
||||
for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) {
|
||||
int64_t bos, seq_len;
|
||||
if (cu_seqlens != nullptr) {
|
||||
bos = cu_seqlens[seq_idx];
|
||||
seq_len = cu_seqlens[seq_idx + 1] - bos;
|
||||
} else {
|
||||
bos = seq_idx;
|
||||
seq_len = 1;
|
||||
}
|
||||
|
||||
int64_t state_read_idx = (state_batch_indices != nullptr)
|
||||
? state_batch_indices[seq_idx]
|
||||
: seq_idx;
|
||||
if (state_read_idx == null_block_id) continue;
|
||||
|
||||
int64_t state_write_idx = (num_accepted_tokens == nullptr)
|
||||
? ((dst_state_batch_indices != nullptr)
|
||||
? dst_state_batch_indices[seq_idx]
|
||||
: state_read_idx)
|
||||
: -1;
|
||||
|
||||
state_t* s = state_ptr + state_read_idx * stride_state_n;
|
||||
|
||||
for (int64_t t = 0; t < seq_len; ++t) {
|
||||
int64_t token_idx = bos + t;
|
||||
const input_t* x_tok = x_ptr + token_idx * stride_x_n;
|
||||
// dt: (N, nheads) — one float per head per token
|
||||
const float* dt_tok = dt_ptr + token_idx * stride_dt_n;
|
||||
const input_t* B_tok = B_ptr + token_idx * stride_BC_n;
|
||||
const input_t* C_tok = C_ptr + token_idx * stride_BC_n;
|
||||
out_t* out_tok = out_ptr + token_idx * stride_out_n;
|
||||
|
||||
#pragma omp parallel for
|
||||
for (int64_t h = 0; h < nheads; ++h) {
|
||||
int64_t g = h / nheads_per_group;
|
||||
const input_t* x_h = x_tok + h * stride_x_h;
|
||||
const input_t* B_g = B_tok + g * stride_BC_g;
|
||||
const input_t* C_g = C_tok + g * stride_BC_g;
|
||||
out_t* out_h = out_tok + h * stride_out_h;
|
||||
state_t* s_h = s + h * stride_state_h;
|
||||
|
||||
// Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing
|
||||
float dt_val = dt_tok[h];
|
||||
if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h];
|
||||
if (dt_softplus) {
|
||||
dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val;
|
||||
}
|
||||
const float A_val = A_ptr[h]; // scalar: same for all dim, dstate
|
||||
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
|
||||
|
||||
const input_t* z_h =
|
||||
(z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h
|
||||
: nullptr;
|
||||
|
||||
vec_op::FP32Vec8 dt_vec(dt_val);
|
||||
// dA = exp(A * dt): A and dt are SCALARS per head, so compute once
|
||||
// and broadcast. This saves 7 redundant std::exp() calls that
|
||||
// FP32Vec8::exp() would otherwise make on the broadcast vector.
|
||||
const float dA_scalar = std::exp(A_val * dt_val);
|
||||
vec_op::FP32Vec8 dA(dA_scalar); // broadcast
|
||||
|
||||
for (int64_t d = 0; d < dim; ++d) {
|
||||
float x_val = static_cast<float>(x_h[d]);
|
||||
|
||||
vec_op::FP32Vec8 out_vec(0.0f);
|
||||
state_t* s_hd = s_h + d * stride_state_d;
|
||||
const input_t* B_g_base = B_g;
|
||||
const input_t* C_g_base = C_g;
|
||||
|
||||
vec_op::FP32Vec8 x_vec(x_val);
|
||||
// dBx = B * x * dt — same dA for all dstate (A is scalar)
|
||||
// s_new = s * dA + B * x * dt
|
||||
|
||||
int64_t n = 0;
|
||||
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
|
||||
vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n)));
|
||||
vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n)));
|
||||
vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n)));
|
||||
|
||||
vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec;
|
||||
vec_op::FP32Vec8 s_new = s_v * dA + dBx;
|
||||
|
||||
state_vec_t(s_new).save(s_hd + n);
|
||||
out_vec = out_vec + s_new * C_v;
|
||||
}
|
||||
|
||||
float out_val = out_vec.reduce_sum();
|
||||
for (; n < dstate; ++n) {
|
||||
// Reuse dA_scalar computed once per head — no exp() re-call
|
||||
float dBx = static_cast<float>(B_g[n]) * x_val * dt_val;
|
||||
float s_new = static_cast<float>(s_hd[n]) * dA_scalar + dBx;
|
||||
s_hd[n] = static_cast<state_t>(s_new);
|
||||
out_val += s_new * static_cast<float>(C_g[n]);
|
||||
}
|
||||
|
||||
if (D_ptr != nullptr) out_val += x_val * D_val;
|
||||
if (z_h != nullptr) {
|
||||
float z_val = static_cast<float>(z_h[d]);
|
||||
float sigmoid = (z_val >= 0)
|
||||
? 1.0f / (1.0f + std::exp(-z_val))
|
||||
: std::exp(z_val) / (1.0f + std::exp(z_val));
|
||||
out_val *= z_val * sigmoid;
|
||||
}
|
||||
out_h[d] = static_cast<out_t>(out_val);
|
||||
}
|
||||
}
|
||||
|
||||
if (num_accepted_tokens != nullptr &&
|
||||
dst_state_batch_indices != nullptr) {
|
||||
int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t];
|
||||
if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) {
|
||||
state_t* dst_s = state_ptr + token_dst_idx * stride_state_n;
|
||||
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (num_accepted_tokens == nullptr && state_write_idx != null_block_id &&
|
||||
state_write_idx != state_read_idx) {
|
||||
state_t* dst_s = state_ptr + state_write_idx * stride_state_n;
|
||||
std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// mamba_chunk_scan_fwd
|
||||
//
|
||||
// Prefill SSM recurrence for Mamba2 / SSD models.
|
||||
//
|
||||
// Key difference from selective_state_update_kernel (decode path):
|
||||
// - #pragma omp parallel for collapse(2) is OUTSIDE the time loop.
|
||||
// Each thread owns a (batch, head) slice and runs the entire token
|
||||
// sequence without any per-token OpenMP synchronisation overhead.
|
||||
// For seqlen=256, this eliminates 256 thread-barrier launches per batch.
|
||||
//
|
||||
// `dt` arrives already processed (float32, after bias + softplus + clamp)
|
||||
// to keep this kernel simple. Preprocessing is done in the Python wrapper.
|
||||
//
|
||||
// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output
|
||||
// tensor, pre-initialised by the caller (zero or from initial_states).
|
||||
// Each (b, h) slice is private to exactly one thread via collapse(2), so
|
||||
// there are no write conflicts.
|
||||
//
|
||||
// D is treated as a scalar per head ([nheads] float32).
|
||||
// ---------------------------------------------------------------------------
|
||||
template <typename input_t>
|
||||
inline void mamba_chunk_scan_fwd_kernel(
|
||||
float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32
|
||||
const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim]
|
||||
const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed)
|
||||
const float* __restrict__ A_ptr, // [nheads] f32
|
||||
const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate]
|
||||
const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate]
|
||||
const float* __restrict__ D_ptr, // [nheads] f32 (nullable)
|
||||
const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable)
|
||||
input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim]
|
||||
const int32_t* __restrict__ cu_seqlens, // [batch+1] int32
|
||||
int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim,
|
||||
int64_t dstate) {
|
||||
using input_vec_t = vec_op::vec_t<input_t>;
|
||||
constexpr int VEC_ELEM_NUM = 8;
|
||||
|
||||
const int64_t nheads_per_group = nheads / ngroups;
|
||||
// states layout: [batch, nheads, headdim, dstate] contiguous (caller
|
||||
// guarantee)
|
||||
const int64_t stride_s_b = nheads * headdim * dstate;
|
||||
const int64_t stride_s_h = headdim * dstate;
|
||||
// stride_s_d = dstate, stride_s_n = 1
|
||||
|
||||
#pragma omp parallel for collapse(2) schedule(static)
|
||||
for (int64_t b = 0; b < batch; ++b) {
|
||||
for (int64_t h = 0; h < nheads; ++h) {
|
||||
const int64_t seq_start = cu_seqlens[b];
|
||||
const int64_t seq_end = cu_seqlens[b + 1];
|
||||
const int64_t g = h / nheads_per_group;
|
||||
|
||||
const float A_val = A_ptr[h];
|
||||
const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f;
|
||||
|
||||
// Working state slice: states[b, h, :, :] — float32, headdim * dstate.
|
||||
// Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB).
|
||||
float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h;
|
||||
|
||||
for (int64_t t = seq_start; t < seq_end; ++t) {
|
||||
const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim;
|
||||
const float* dt_h = dt_ptr + t * nheads + h;
|
||||
const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate;
|
||||
const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate;
|
||||
const input_t* z_h = (z_ptr != nullptr)
|
||||
? z_ptr + t * nheads * headdim + h * headdim
|
||||
: nullptr;
|
||||
input_t* out_h = out_ptr + t * nheads * headdim + h * headdim;
|
||||
|
||||
const float dt_val = *dt_h;
|
||||
const float dA_val = std::exp(A_val * dt_val);
|
||||
const vec_op::FP32Vec8 dA_vec(dA_val); // broadcast scalar
|
||||
const vec_op::FP32Vec8 dt_vec(dt_val);
|
||||
|
||||
for (int64_t d = 0; d < headdim; ++d) {
|
||||
const float x_val = static_cast<float>(x_h[d]);
|
||||
float* s_bhd = s_bh + d * dstate; // [dstate] contiguous float32
|
||||
|
||||
// Vectorised SSM update + readout over dstate:
|
||||
// s_new = s * dA + x * dt * B
|
||||
// y += s_new * C
|
||||
int64_t n = 0;
|
||||
vec_op::FP32Vec8 y_vec(0.0f);
|
||||
const vec_op::FP32Vec8 x_vec(x_val);
|
||||
|
||||
for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) {
|
||||
const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n)));
|
||||
const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n)));
|
||||
const vec_op::FP32Vec8 s_v(s_bhd + n);
|
||||
|
||||
const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v;
|
||||
s_new.save(s_bhd + n);
|
||||
y_vec = y_vec + s_new * C_v;
|
||||
}
|
||||
|
||||
float y_val = y_vec.reduce_sum();
|
||||
|
||||
// Scalar tail for remaining dstate elements
|
||||
for (; n < dstate; ++n) {
|
||||
const float B_n = static_cast<float>(B_g[n]);
|
||||
const float C_n = static_cast<float>(C_g[n]);
|
||||
const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n;
|
||||
s_bhd[n] = s_new;
|
||||
y_val += s_new * C_n;
|
||||
}
|
||||
|
||||
// D skip connection (scalar per head)
|
||||
if (D_ptr != nullptr) y_val += x_val * D_val;
|
||||
|
||||
// z gating: out = y * z * sigmoid(z) (SiLU)
|
||||
if (z_h != nullptr) {
|
||||
const float z_val = static_cast<float>(z_h[d]);
|
||||
const float sigmoid =
|
||||
(z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val))
|
||||
: std::exp(z_val) / (1.0f + std::exp(z_val));
|
||||
y_val *= z_val * sigmoid;
|
||||
}
|
||||
|
||||
out_h[d] = static_cast<input_t>(y_val);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mamba_cpu
|
||||
@@ -213,6 +213,32 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
|
||||
torch::Tensor slot_mapping,
|
||||
const int64_t block_size);
|
||||
|
||||
at::Tensor causal_conv1d_update_cpu_impl(
|
||||
at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight,
|
||||
const c10::optional<at::Tensor>& bias,
|
||||
const c10::optional<std::string>& activation,
|
||||
const c10::optional<at::Tensor>& conv_state_indices,
|
||||
const c10::optional<at::Tensor>& query_start_loc, int64_t pad_slot_id);
|
||||
|
||||
void selective_state_update_cpu_impl(
|
||||
at::Tensor& state, const at::Tensor& x, const at::Tensor& dt,
|
||||
const at::Tensor& A, const at::Tensor& B, const at::Tensor& C,
|
||||
const c10::optional<at::Tensor>& D, const c10::optional<at::Tensor>& z,
|
||||
const c10::optional<at::Tensor>& dt_bias, bool dt_softplus,
|
||||
const c10::optional<at::Tensor>& state_batch_indices,
|
||||
const c10::optional<at::Tensor>& dst_state_batch_indices,
|
||||
int64_t null_block_id, at::Tensor& out,
|
||||
const c10::optional<at::Tensor>& num_accepted_tokens,
|
||||
const c10::optional<at::Tensor>& cu_seqlens);
|
||||
|
||||
void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states,
|
||||
const at::Tensor& x, const at::Tensor& dt,
|
||||
const at::Tensor& A, const at::Tensor& B,
|
||||
const at::Tensor& C,
|
||||
const c10::optional<at::Tensor>& D,
|
||||
const c10::optional<at::Tensor>& z,
|
||||
const at::Tensor& cu_seqlens);
|
||||
|
||||
void init_cpu_memory_env(std::vector<int64_t> node_ids);
|
||||
|
||||
namespace cpu_utils {
|
||||
@@ -595,6 +621,30 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
"block_size) -> ()",
|
||||
&compute_slot_mapping_kernel_impl);
|
||||
|
||||
// Mamba CPU kernels
|
||||
ops.def(
|
||||
"causal_conv1d_update_cpu_vec("
|
||||
"Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, "
|
||||
"Tensor? bias, str? activation, Tensor? conv_state_indices, "
|
||||
"Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor",
|
||||
&causal_conv1d_update_cpu_impl);
|
||||
|
||||
ops.def(
|
||||
"selective_state_update_cpu("
|
||||
"Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
|
||||
"Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, "
|
||||
"Tensor? state_batch_indices, Tensor? dst_state_batch_indices, "
|
||||
"SymInt null_block_id, Tensor(a13!) out, "
|
||||
"Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()",
|
||||
&selective_state_update_cpu_impl);
|
||||
|
||||
ops.def(
|
||||
"mamba_chunk_scan_fwd_cpu("
|
||||
"Tensor(a0!) out, Tensor(a1!) final_states, "
|
||||
"Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, "
|
||||
"Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()",
|
||||
&mamba_chunk_scan_fwd_cpu_impl);
|
||||
|
||||
ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
|
||||
|
||||
// Speculative decoding kernels
|
||||
|
||||
@@ -1025,6 +1025,9 @@ __global__ void gather_and_maybe_dequant_cache(
|
||||
batch_offset += offset;
|
||||
int32_t block_table_id = batch_offset / block_size;
|
||||
int32_t slot_id = batch_offset % block_size;
|
||||
// seq_starts may push the block index past the end of the batch's block
|
||||
// table row.
|
||||
if (block_table_id >= block_table_stride) continue;
|
||||
int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
|
||||
int32_t block_id = block_table[block_table_offset];
|
||||
int64_t cache_offset =
|
||||
@@ -1174,7 +1177,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
const int32_t num_reqs, const int32_t block_size,
|
||||
const int32_t total_tokens, const int64_t block_table_stride,
|
||||
const int64_t cache_block_stride, const int64_t cache_entry_stride,
|
||||
const int64_t dst_entry_stride) {
|
||||
const int64_t dst_entry_stride,
|
||||
const int32_t* __restrict__ seq_starts) { // Optional source offsets
|
||||
const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
|
||||
if (flat_warp_id >= total_tokens) return;
|
||||
const int lane_id = threadIdx.x & 31;
|
||||
@@ -1192,7 +1196,8 @@ __global__ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
|
||||
// Compute physical token address via block table
|
||||
const int out_token_id = flat_warp_id;
|
||||
const int token_offset = out_token_id - workspace_starts[req_id];
|
||||
int token_offset = out_token_id - workspace_starts[req_id];
|
||||
if (seq_starts != nullptr) token_offset += seq_starts[req_id];
|
||||
const int cache_block_idx = token_offset / block_size;
|
||||
const int offset_in_block = token_offset % block_size;
|
||||
const int physical_block =
|
||||
@@ -1383,9 +1388,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]
|
||||
torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
|
||||
torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
|
||||
torch::stable::Tensor const& seq_lens, // [BATCH]
|
||||
torch::stable::Tensor const& workspace_starts, // [BATCH]
|
||||
int64_t batch_size) {
|
||||
int64_t batch_size,
|
||||
std::optional<torch::stable::Tensor> seq_starts = std::nullopt) {
|
||||
torch::stable::accelerator::DeviceGuard device_guard(
|
||||
src_cache.get_device_index());
|
||||
const cudaStream_t stream = get_current_cuda_stream();
|
||||
@@ -1396,20 +1401,25 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
STD_TORCH_CHECK(
|
||||
block_table.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"block_table must be int32");
|
||||
STD_TORCH_CHECK(seq_lens.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"seq_lens must be int32");
|
||||
STD_TORCH_CHECK(
|
||||
workspace_starts.scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"workspace_starts must be int32");
|
||||
if (seq_starts.has_value()) {
|
||||
STD_TORCH_CHECK(
|
||||
seq_starts.value().scalar_type() == torch::headeronly::ScalarType::Int,
|
||||
"seq_starts must be int32");
|
||||
}
|
||||
|
||||
STD_TORCH_CHECK(src_cache.device() == dst.device(),
|
||||
"src_cache and dst must be on the same device");
|
||||
STD_TORCH_CHECK(src_cache.device() == block_table.device(),
|
||||
"src_cache and block_table must be on the same device");
|
||||
STD_TORCH_CHECK(src_cache.device() == seq_lens.device(),
|
||||
"src_cache and seq_lens must be on the same device");
|
||||
STD_TORCH_CHECK(src_cache.device() == workspace_starts.device(),
|
||||
"src_cache and workspace_starts must be on the same device");
|
||||
if (seq_starts.has_value()) {
|
||||
STD_TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
|
||||
"src_cache and seq_starts must be on the same device");
|
||||
}
|
||||
auto dtype = src_cache.scalar_type();
|
||||
STD_TORCH_CHECK(
|
||||
dtype == torch::headeronly::ScalarType::Byte || // uint8
|
||||
@@ -1438,6 +1448,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
constexpr int warps_per_block = 8;
|
||||
const int grid_size = (total_tokens + warps_per_block - 1) / warps_per_block;
|
||||
const int block_size_threads = warps_per_block * 32; // 256 threads
|
||||
const int32_t* seq_starts_ptr =
|
||||
seq_starts.has_value() ? seq_starts.value().const_data_ptr<int32_t>()
|
||||
: nullptr;
|
||||
|
||||
vllm::cp_gather_and_upconvert_fp8_kv_cache<<<grid_size, block_size_threads, 0,
|
||||
stream>>>(
|
||||
@@ -1446,7 +1459,7 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
workspace_starts.const_data_ptr<int32_t>(),
|
||||
static_cast<int32_t>(batch_size), block_size, total_tokens,
|
||||
block_table_stride, cache_block_stride, cache_entry_stride,
|
||||
dst_entry_stride);
|
||||
dst_entry_stride, seq_starts_ptr);
|
||||
}
|
||||
|
||||
// Macro to dispatch the kernel based on the data type.
|
||||
|
||||
@@ -9,14 +9,16 @@ void topk_softmax(torch::stable::Tensor& topk_weights,
|
||||
torch::stable::Tensor& topk_indices,
|
||||
torch::stable::Tensor& token_expert_indices,
|
||||
torch::stable::Tensor& gating_output, bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias);
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
std::optional<torch::stable::Tensor> is_padding);
|
||||
|
||||
void topk_sigmoid(torch::stable::Tensor& topk_weights,
|
||||
torch::stable::Tensor& topk_indices,
|
||||
torch::stable::Tensor& token_expert_indices,
|
||||
torch::stable::Tensor& gating_output, bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor);
|
||||
double routed_scaling_factor,
|
||||
std::optional<torch::stable::Tensor> is_padding);
|
||||
|
||||
void topk_softplus_sqrt(
|
||||
torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
|
||||
@@ -25,7 +27,8 @@ void topk_softplus_sqrt(
|
||||
double routed_scaling_factor,
|
||||
const std::optional<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid);
|
||||
const std::optional<torch::stable::Tensor>& tid2eid,
|
||||
const std::optional<torch::stable::Tensor>& is_padding);
|
||||
|
||||
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
|
||||
std::optional<torch::stable::Tensor> topk_ids,
|
||||
|
||||
@@ -174,7 +174,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
|
||||
const int end_expert,
|
||||
const bool renormalize,
|
||||
const float* bias,
|
||||
const double routed_scaling_factor)
|
||||
const double routed_scaling_factor,
|
||||
const bool* is_padding)
|
||||
{
|
||||
|
||||
using cub_kvp = cub::KeyValuePair<int, float>;
|
||||
@@ -228,12 +229,14 @@ __launch_bounds__(TPB) __global__ void moeTopK(
|
||||
const int expert = result_kvp.key;
|
||||
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
|
||||
const bool should_process_row = row_is_active && node_uses_expert;
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[block_row];
|
||||
|
||||
const int idx = k * block_row + k_idx;
|
||||
// Return the unbiased scores for output weights
|
||||
output[idx] = inputs_after_softmax[thread_read_offset + expert];
|
||||
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
|
||||
assert(indices[idx] >= 0);
|
||||
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
|
||||
: (should_process_row ? (expert - start_expert) : num_experts);
|
||||
assert(is_pad_row || indices[idx] >= 0);
|
||||
source_rows[idx] = k_idx * num_rows + block_row;
|
||||
if (renormalize) {
|
||||
selected_sum += inputs_after_softmax[thread_read_offset + expert];
|
||||
@@ -277,7 +280,7 @@ template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WA
|
||||
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
|
||||
int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
|
||||
const float* bias, const double routed_scaling_factor)
|
||||
const float* bias, const double routed_scaling_factor, const bool* is_padding)
|
||||
{
|
||||
static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>,
|
||||
@@ -545,12 +548,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
// Add a guard to ignore experts not included by this node
|
||||
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
|
||||
const bool should_process_row = row_is_active && node_uses_expert;
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
|
||||
|
||||
// The lead thread from each sub-group will write out the final results to global memory. (This will be a
|
||||
// single) thread per row of the input/output matrices.
|
||||
const int idx = k * thread_row + k_idx;
|
||||
output[idx] = max_val;
|
||||
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
|
||||
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
|
||||
: (should_process_row ? (expert - start_expert) : NUM_EXPERTS);
|
||||
source_rows[idx] = k_idx * num_rows + thread_row;
|
||||
if (renormalize) {
|
||||
selected_sum += max_val;
|
||||
@@ -605,7 +610,7 @@ struct TopkConstants
|
||||
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
|
||||
void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
|
||||
int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
|
||||
const float* bias, const double routed_scaling_factor, cudaStream_t stream)
|
||||
const float* bias, const double routed_scaling_factor, cudaStream_t stream, const bool* is_padding)
|
||||
{
|
||||
static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
|
||||
using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
|
||||
@@ -616,7 +621,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
|
||||
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
|
||||
topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding);
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
@@ -627,7 +632,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
#else
|
||||
#define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
|
||||
if (WARP_SIZE == 64) { \
|
||||
@@ -635,13 +640,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, routed_scaling_factor, stream); \
|
||||
bias, routed_scaling_factor, stream, is_padding); \
|
||||
} else if (WARP_SIZE == 32) { \
|
||||
topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES, \
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, routed_scaling_factor, stream); \
|
||||
bias, routed_scaling_factor, stream, is_padding); \
|
||||
} else { \
|
||||
assert(false && \
|
||||
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
|
||||
@@ -661,7 +666,8 @@ void topkGatingKernelLauncher(
|
||||
const bool renormalize,
|
||||
const float* bias,
|
||||
const double routed_scaling_factor,
|
||||
cudaStream_t stream) {
|
||||
cudaStream_t stream,
|
||||
const bool* is_padding) {
|
||||
static constexpr int WARPS_PER_TB = 4;
|
||||
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
|
||||
#ifndef USE_ROCM
|
||||
@@ -736,7 +742,7 @@ void topkGatingKernelLauncher(
|
||||
}
|
||||
moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
|
||||
workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
|
||||
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);
|
||||
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor, is_padding);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -755,7 +761,8 @@ void dispatch_topk_launch(
|
||||
int num_tokens, int num_experts, int topk, bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor,
|
||||
cudaStream_t stream)
|
||||
cudaStream_t stream,
|
||||
std::optional<torch::stable::Tensor> is_padding)
|
||||
{
|
||||
const float* bias_ptr = nullptr;
|
||||
if (bias.has_value()) {
|
||||
@@ -769,6 +776,18 @@ void dispatch_topk_launch(
|
||||
bias_ptr = bias_tensor.const_data_ptr<float>();
|
||||
}
|
||||
|
||||
const bool* is_padding_ptr = nullptr;
|
||||
if (is_padding.has_value()) {
|
||||
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
|
||||
STD_TORCH_CHECK(is_padding_tensor.scalar_type() == torch::headeronly::ScalarType::Bool,
|
||||
"is_padding tensor must be bool");
|
||||
STD_TORCH_CHECK(is_padding_tensor.dim() == 1, "is_padding tensor must be 1D");
|
||||
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
|
||||
"is_padding size mismatch, expected: ", num_tokens);
|
||||
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), "is_padding tensor must be contiguous");
|
||||
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
|
||||
}
|
||||
|
||||
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
|
||||
vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
|
||||
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
|
||||
@@ -777,7 +796,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
|
||||
} else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
|
||||
vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
|
||||
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
|
||||
@@ -786,7 +805,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
|
||||
} else {
|
||||
STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
|
||||
vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
|
||||
@@ -796,7 +815,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -806,7 +825,8 @@ void topk_softmax(
|
||||
torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
|
||||
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
|
||||
bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias)
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
std::optional<torch::stable::Tensor> is_padding)
|
||||
{
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
@@ -825,15 +845,15 @@ void topk_softmax(
|
||||
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
|
||||
dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, 1.0, stream);
|
||||
bias, 1.0, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
|
||||
dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, 1.0, stream);
|
||||
bias, 1.0, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, 1.0, stream);
|
||||
bias, 1.0, stream, is_padding);
|
||||
} else {
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
|
||||
}
|
||||
@@ -846,7 +866,8 @@ void topk_sigmoid(
|
||||
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
|
||||
bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor)
|
||||
double routed_scaling_factor,
|
||||
std::optional<torch::stable::Tensor> is_padding)
|
||||
{
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
@@ -865,15 +886,15 @@ void topk_sigmoid(
|
||||
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
|
||||
dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
|
||||
dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, routed_scaling_factor, stream);
|
||||
bias, routed_scaling_factor, stream, is_padding);
|
||||
} else {
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
|
||||
}
|
||||
|
||||
@@ -71,6 +71,80 @@ __device__ __forceinline__ float toFloat(T value) {
|
||||
}
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// Adapted from:
|
||||
// https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh
|
||||
template <typename OutIndType, typename HashIndType>
|
||||
__launch_bounds__(128) __global__
|
||||
void dsv4HashTopkSoftplusSqrt(const float* input, float* output,
|
||||
OutIndType* indices, int num_rows,
|
||||
int num_experts, float routed_scaling_factor,
|
||||
const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid,
|
||||
const bool* is_padding) {
|
||||
const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
|
||||
const int lane = threadIdx.x % 32;
|
||||
if (warp >= num_rows) return;
|
||||
const int64_t token_id = load_index_as_int64(input_ids, warp);
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[warp];
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
int expert = 0;
|
||||
float weight = 0.f;
|
||||
if (lane < 6 && !is_pad_row) {
|
||||
// only load and calculate for 6 experts
|
||||
expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
|
||||
const float x = input[warp * num_experts + expert];
|
||||
weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
|
||||
if (isnan(weight)) {
|
||||
weight = 0.f;
|
||||
}
|
||||
}
|
||||
float weight_sum = weight;
|
||||
#pragma unroll
|
||||
for (int mask = 16; mask > 0; mask >>= 1) {
|
||||
// sum in warp
|
||||
weight_sum += VLLM_SHFL_XOR_SYNC(weight_sum, mask);
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
|
||||
cudaTriggerProgrammaticLaunchCompletion();
|
||||
#endif
|
||||
if (lane < 6) {
|
||||
const int offset = warp * 6 + lane;
|
||||
output[offset] =
|
||||
weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);
|
||||
indices[offset] = !is_pad_row ? static_cast<OutIndType>(expert)
|
||||
: static_cast<OutIndType>(-1);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename OutIndType, typename HashIndType>
|
||||
void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
|
||||
int num_rows, int num_experts,
|
||||
double routed_scaling_factor,
|
||||
const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid, cudaStream_t stream,
|
||||
const bool* is_padding) {
|
||||
if (num_rows == 0) return;
|
||||
auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
|
||||
cudaLaunchConfig_t config = {};
|
||||
config.gridDim = (num_rows + 3) / 4;
|
||||
config.blockDim = 128;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attr;
|
||||
attr.id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attr.val.programmaticStreamSerializationAllowed = 1;
|
||||
config.attrs = &attr;
|
||||
config.numAttrs = 1;
|
||||
const float scale = static_cast<float>(routed_scaling_factor);
|
||||
cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
|
||||
num_experts, scale, input_ids, tid2eid, is_padding);
|
||||
}
|
||||
#endif
|
||||
|
||||
// ====================== TopK softplus_sqrt things
|
||||
// ===============================
|
||||
|
||||
@@ -99,7 +173,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
const int num_rows, IndType* indices, int* source_rows, const int k,
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const HashIndType* input_ids, const HashIndType* tid2eid) {
|
||||
const HashIndType* input_ids, const HashIndType* tid2eid,
|
||||
const bool* is_padding) {
|
||||
static_assert(std::is_same_v<InputType, float> ||
|
||||
std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>,
|
||||
@@ -164,6 +239,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
return;
|
||||
}
|
||||
const bool row_is_active = finished ? !finished[thread_row] : true;
|
||||
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
|
||||
|
||||
// We finally start setting up the read pointers for each thread. First, each
|
||||
// thread jumps to the start of the row it will read.
|
||||
@@ -182,9 +258,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
cudaGridDependencySynchronize();
|
||||
#endif
|
||||
|
||||
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
|
||||
// to float
|
||||
if constexpr (std::is_same_v<InputType, float>) {
|
||||
if (is_pad_row) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
row_chunk[ii] = 0.f;
|
||||
}
|
||||
} else if constexpr (std::is_same_v<InputType, float>) {
|
||||
using VecType = AlignedArray<float, ELTS_PER_LDG>;
|
||||
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
|
||||
const VecType* vec_thread_read_ptr =
|
||||
@@ -248,12 +327,22 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
if constexpr (USE_HASH) {
|
||||
const int64_t token_id = load_index_as_int64(input_ids, thread_row);
|
||||
const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);
|
||||
if (!is_pad_row) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
row_chunk[ii] = sqrtf(val);
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
val = sqrtf(val);
|
||||
|
||||
// Dummy/padding tokens can result in NaN values, so
|
||||
// clamp them to 0.0. Note: this clamp could likely be removed if
|
||||
// 'is_padding' is made mandatory
|
||||
if (isnan(val)) {
|
||||
val = 0.f;
|
||||
}
|
||||
row_chunk[ii] = val;
|
||||
}
|
||||
}
|
||||
float selected_sum = 0.f;
|
||||
#pragma unroll
|
||||
@@ -268,7 +357,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
if (expert == expert_idx) {
|
||||
indices[idx] = static_cast<IndType>(expert);
|
||||
indices[idx] = !is_pad_row ? static_cast<IndType>(expert)
|
||||
: static_cast<IndType>(-1);
|
||||
selected_sum += row_chunk[ii];
|
||||
break;
|
||||
}
|
||||
@@ -312,23 +402,31 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
#endif
|
||||
return;
|
||||
} else {
|
||||
if (!is_pad_row) {
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
// Compute softplus: log(1 + exp(val)) with numerical stability
|
||||
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
val = sqrtf(val);
|
||||
if (correction_bias) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
val = val + correction_bias[expert_idx];
|
||||
for (int ii = 0; ii < VPT; ++ii) {
|
||||
float val = row_chunk[ii];
|
||||
float val_b = val * beta;
|
||||
// Compute softplus: log(1 + exp(val)) with numerical stability
|
||||
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
|
||||
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
|
||||
val = sqrtf(val);
|
||||
// Dummy/padding tokens can result in NaN values, so
|
||||
// clamp them to 0.0. Note: this clamp could likely be removed if
|
||||
// 'is_padding' is made mandatory
|
||||
if (isnan(val)) {
|
||||
val = 0.f;
|
||||
}
|
||||
if (correction_bias) {
|
||||
const int group_id = ii / ELTS_PER_LDG;
|
||||
const int local_id = ii % ELTS_PER_LDG;
|
||||
const int expert_idx = first_elt_read_by_thread +
|
||||
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
|
||||
local_id;
|
||||
val = val + correction_bias[expert_idx];
|
||||
}
|
||||
row_chunk[ii] = val;
|
||||
}
|
||||
row_chunk[ii] = val;
|
||||
}
|
||||
|
||||
// Original TopK path: find top-k experts by score
|
||||
@@ -383,18 +481,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
// Add a guard to ignore experts not included by this node
|
||||
const bool node_uses_expert =
|
||||
expert >= start_expert && expert < end_expert;
|
||||
const bool should_process_row = row_is_active && node_uses_expert;
|
||||
const bool should_process_row =
|
||||
row_is_active && node_uses_expert && !is_pad_row;
|
||||
|
||||
// The lead thread from each sub-group will write out the final results
|
||||
// to global memory. (This will be a single) thread per row of the
|
||||
// input/output matrices.
|
||||
const int idx = k * thread_row + k_idx;
|
||||
if (correction_bias != nullptr) {
|
||||
if (correction_bias != nullptr && should_process_row) {
|
||||
max_val -= correction_bias[expert];
|
||||
}
|
||||
output[idx] = max_val;
|
||||
indices[idx] =
|
||||
should_process_row ? (expert - start_expert) : NUM_EXPERTS;
|
||||
!is_pad_row ? expert - start_expert : static_cast<IndType>(-1);
|
||||
source_rows[idx] = k_idx * num_rows + thread_row;
|
||||
if (renormalize) {
|
||||
selected_sum += max_val;
|
||||
@@ -477,7 +576,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
const int start_expert, const int end_expert, const bool renormalize,
|
||||
double routed_scaling_factor, const float* correction_bias,
|
||||
const bool use_hash, const HashIndType* input_ids,
|
||||
const HashIndType* tid2eid, cudaStream_t stream) {
|
||||
const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) {
|
||||
static constexpr int BYTES_PER_LDG =
|
||||
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
|
||||
using Constants =
|
||||
@@ -506,12 +605,12 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
|
||||
indices, source_row, k, start_expert, end_expert,
|
||||
renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid);
|
||||
input_ids, tid2eid, is_padding);
|
||||
#else
|
||||
kernel<<<num_blocks, block_dim, 0, stream>>>(
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert,
|
||||
end_expert, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid);
|
||||
input_ids, tid2eid, is_padding);
|
||||
#endif
|
||||
})
|
||||
}
|
||||
@@ -525,7 +624,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
|
||||
stream);
|
||||
stream, is_padding);
|
||||
#else
|
||||
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
|
||||
if (WARP_SIZE == 64) { \
|
||||
@@ -534,14 +633,14 @@ void topkGatingSoftplusSqrtLauncherHelper(
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, \
|
||||
tid2eid, stream); \
|
||||
tid2eid, stream, is_padding); \
|
||||
} else if (WARP_SIZE == 32) { \
|
||||
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
|
||||
MAX_BYTES>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
routed_scaling_factor, correction_bias, use_hash, input_ids, \
|
||||
tid2eid, stream); \
|
||||
tid2eid, stream, is_padding); \
|
||||
} else { \
|
||||
assert(false && \
|
||||
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
|
||||
@@ -555,7 +654,18 @@ void topkGatingSoftplusSqrtKernelLauncher(
|
||||
const int topk, const bool renormalize, double routed_scaling_factor,
|
||||
const float* correction_bias, const bool use_hash,
|
||||
const HashIndType* input_ids, const HashIndType* tid2eid,
|
||||
cudaStream_t stream) {
|
||||
cudaStream_t stream, const bool* is_padding) {
|
||||
#ifndef USE_ROCM
|
||||
if constexpr (std::is_same_v<InputType, float>) {
|
||||
if (use_hash && topk == 6 && renormalize &&
|
||||
(num_experts == 256 || num_experts == 384)) {
|
||||
launchDsv4HashTopk<IndType, HashIndType>(
|
||||
gating_output, topk_weights, topk_indices, num_tokens, num_experts,
|
||||
routed_scaling_factor, input_ids, tid2eid, stream, is_padding);
|
||||
return;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
static constexpr int WARPS_PER_TB = 4;
|
||||
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
|
||||
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
|
||||
@@ -650,7 +760,8 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
int num_experts, int topk, bool renormalize, double routed_scaling_factor,
|
||||
const std::optional<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
|
||||
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,
|
||||
const std::optional<torch::stable::Tensor>& is_padding) {
|
||||
const float* bias_ptr = nullptr;
|
||||
if (correction_bias.has_value()) {
|
||||
bias_ptr = correction_bias.value().const_data_ptr<float>();
|
||||
@@ -659,6 +770,22 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
auto launch = [&](auto* topk_indices_ptr) {
|
||||
using OutIndType =
|
||||
typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
|
||||
|
||||
const bool* is_padding_ptr = nullptr;
|
||||
if (is_padding.has_value()) {
|
||||
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
|
||||
STD_TORCH_CHECK(is_padding_tensor.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Bool,
|
||||
"is_padding tensor must be bool");
|
||||
STD_TORCH_CHECK(is_padding_tensor.dim() == 1,
|
||||
"is_padding tensor must be 1D");
|
||||
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
|
||||
"is_padding size mismatch, expected: ", num_tokens);
|
||||
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(),
|
||||
"is_padding tensor must be contiguous");
|
||||
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
|
||||
}
|
||||
|
||||
if (tid2eid.has_value()) {
|
||||
STD_TORCH_CHECK(input_ids.has_value(),
|
||||
"input_ids is required for hash MoE");
|
||||
@@ -673,7 +800,7 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
|
||||
tid2eid.value().const_data_ptr<int64_t>(), stream);
|
||||
tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr);
|
||||
} else {
|
||||
STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
|
||||
torch::headeronly::ScalarType::Int);
|
||||
@@ -683,7 +810,7 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, true, input_ids.value().const_data_ptr<int>(),
|
||||
tid2eid.value().const_data_ptr<int>(), stream);
|
||||
tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);
|
||||
}
|
||||
} else {
|
||||
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
|
||||
@@ -691,7 +818,7 @@ void dispatch_topk_softplus_sqrt_launch(
|
||||
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
|
||||
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
|
||||
bias_ptr, false, static_cast<const OutIndType*>(nullptr),
|
||||
static_cast<const OutIndType*>(nullptr), stream);
|
||||
static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -715,7 +842,8 @@ void topk_softplus_sqrt(
|
||||
bool renormalize, double routed_scaling_factor,
|
||||
const std::optional<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& tid2eid) {
|
||||
const std::optional<torch::stable::Tensor>& tid2eid,
|
||||
const std::optional<torch::stable::Tensor>& is_padding) {
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
const int topk = topk_weights.size(-1);
|
||||
@@ -728,21 +856,22 @@ void topk_softplus_sqrt(
|
||||
dispatch_topk_softplus_sqrt_launch<float>(
|
||||
gating_output.const_data_ptr<float>(), topk_weights, topk_indices,
|
||||
token_expert_indices, num_tokens, num_experts, topk, renormalize,
|
||||
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
|
||||
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream,
|
||||
is_padding);
|
||||
} else if (gating_output.scalar_type() ==
|
||||
torch::headeronly::ScalarType::Half) {
|
||||
dispatch_topk_softplus_sqrt_launch<__half>(
|
||||
reinterpret_cast<const __half*>(gating_output.const_data_ptr()),
|
||||
topk_weights, topk_indices, token_expert_indices, num_tokens,
|
||||
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid, stream);
|
||||
input_ids, tid2eid, stream, is_padding);
|
||||
} else if (gating_output.scalar_type() ==
|
||||
torch::headeronly::ScalarType::BFloat16) {
|
||||
dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),
|
||||
topk_weights, topk_indices, token_expert_indices, num_tokens,
|
||||
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
|
||||
input_ids, tid2eid, stream);
|
||||
input_ids, tid2eid, stream, is_padding);
|
||||
} else {
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
|
||||
gating_output.scalar_type());
|
||||
|
||||
@@ -8,19 +8,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
|
||||
m.def(
|
||||
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
|
||||
"bias) -> ()");
|
||||
"bias, Tensor? is_padding) -> ()");
|
||||
|
||||
// Apply topk sigmoid to the gating outputs.
|
||||
m.def(
|
||||
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, "
|
||||
"Tensor? bias, float routed_scaling_factor) -> ()");
|
||||
"Tensor? bias, float routed_scaling_factor, Tensor? is_padding) -> ()");
|
||||
|
||||
m.def(
|
||||
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, float "
|
||||
"routed_scaling_factor, Tensor? "
|
||||
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
|
||||
"bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()");
|
||||
|
||||
// Calculate the result of moe by summing up the partial results
|
||||
// from all selected experts. topk_ids/expert_map are optional and, when
|
||||
|
||||
@@ -527,9 +527,9 @@ void cp_gather_and_upconvert_fp8_kv_cache(
|
||||
// 656]
|
||||
torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
|
||||
torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
|
||||
torch::stable::Tensor const& seq_lens, // [BATCH]
|
||||
torch::stable::Tensor const& workspace_starts, // [BATCH]
|
||||
int64_t batch_size);
|
||||
int64_t batch_size,
|
||||
std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
|
||||
|
||||
// Indexer K quantization and cache function
|
||||
void indexer_k_quant_and_cache(
|
||||
|
||||
@@ -39,11 +39,15 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
|
||||
// AWQ zeros: (size_k // group_size, size_n // 8)
|
||||
const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
|
||||
int32_t group_size) {
|
||||
int32_t val =
|
||||
qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
|
||||
int32_t zero =
|
||||
qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
|
||||
blockIdx.y];
|
||||
// Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
|
||||
// blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
|
||||
// same row (stride 1) instead of striding across rows (stride size_n/8).
|
||||
int col = blockIdx.y * 32 + threadIdx.x;
|
||||
if (col >= size_n / 8) return;
|
||||
(void)size_k;
|
||||
|
||||
int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
|
||||
int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
|
||||
int32_t new_val = 0;
|
||||
|
||||
#pragma unroll
|
||||
@@ -58,7 +62,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
|
||||
zero >>= 4;
|
||||
}
|
||||
|
||||
output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
|
||||
output[blockIdx.x * (size_n / 8) + col] = new_val;
|
||||
}
|
||||
|
||||
torch::stable::Tensor marlin_int4_fp8_preprocess(
|
||||
@@ -102,7 +106,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess(
|
||||
"qweight.size(0) % qzeros.size(0) != 0");
|
||||
STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
|
||||
|
||||
dim3 blocks(size_k / 32, size_n / 8);
|
||||
dim3 blocks(size_k, (size_n / 8 + 31) / 32);
|
||||
marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
|
||||
reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(output.mutable_data_ptr()),
|
||||
|
||||
@@ -847,8 +847,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C_cache_ops, ops) {
|
||||
|
||||
ops.def(
|
||||
"cp_gather_and_upconvert_fp8_kv_cache(Tensor src_cache, Tensor! dst, "
|
||||
"Tensor block_table, Tensor seq_lens, Tensor workspace_starts, int "
|
||||
"batch_size) -> ()");
|
||||
"Tensor block_table, Tensor workspace_starts, int batch_size, Tensor? "
|
||||
"seq_starts) -> ()");
|
||||
|
||||
ops.def(
|
||||
"indexer_k_quant_and_cache(Tensor k, Tensor! kv_cache, Tensor "
|
||||
|
||||
+2
-2
@@ -22,7 +22,7 @@
|
||||
# docker buildx bake -f docker/docker-bake.hcl -f docker/versions.json
|
||||
# =============================================================================
|
||||
|
||||
ARG CUDA_VERSION=13.0.2
|
||||
ARG CUDA_VERSION=13.0.3
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG UBUNTU_VERSION=22.04
|
||||
|
||||
@@ -793,7 +793,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
|
||||
# https://docs.flashinfer.ai/installation.html
|
||||
# From versions.json: .flashinfer.version
|
||||
ARG FLASHINFER_VERSION=0.6.14
|
||||
ARG FLASHINFER_VERSION=0.6.15.post1
|
||||
RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
|
||||
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
+53
-31
@@ -339,18 +339,17 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
|
||||
|
||||
# RIXL/UCX build stages
|
||||
FROM base AS build_rixl
|
||||
ARG RIXL_BRANCH="39be1de8"
|
||||
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
|
||||
ARG UCX_BRANCH="bfb51733"
|
||||
# NIXL/UCX build stages
|
||||
FROM base AS build_nixl
|
||||
ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
|
||||
ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
|
||||
ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
|
||||
ARG UCX_REPO="https://github.com/openucx/ucx.git"
|
||||
ENV ROCM_PATH=/opt/rocm
|
||||
ENV UCX_HOME=/usr/local/ucx
|
||||
ENV RIXL_HOME=/usr/local/rixl
|
||||
ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
|
||||
ENV NIXL_HOME=/usr/local/nixl
|
||||
|
||||
# RIXL build system dependences and RDMA support
|
||||
# NIXL build system dependencies and RDMA support
|
||||
RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
|
||||
libgrpc-dev \
|
||||
libgrpc++-dev \
|
||||
@@ -368,7 +367,8 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system meson auditwheel patchelf tomlkit
|
||||
uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \
|
||||
auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
cd /usr/local/src && \
|
||||
@@ -396,30 +396,50 @@ ENV PATH=/usr/local/ucx/bin:$PATH
|
||||
ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/ccache \
|
||||
git clone ${RIXL_REPO} /opt/rixl && \
|
||||
cd /opt/rixl && \
|
||||
git checkout ${RIXL_BRANCH} && \
|
||||
git clone ${NIXL_REPO} /opt/nixl && \
|
||||
cd /opt/nixl && \
|
||||
git checkout ${NIXL_BRANCH} && \
|
||||
CC="ccache gcc" CXX="ccache g++" \
|
||||
meson setup build --prefix=${RIXL_HOME} \
|
||||
meson setup build --prefix=${NIXL_HOME} \
|
||||
-Ducx_path=${UCX_HOME} \
|
||||
-Drocm_path=${ROCM_PATH} && \
|
||||
-Dwheel_variant=rocm \
|
||||
-Dbuild_tests=false \
|
||||
-Dbuild_examples=false && \
|
||||
cd build && \
|
||||
ninja -j$(nproc) && \
|
||||
ninja install
|
||||
ninja install && \
|
||||
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
|
||||
> /etc/ld.so.conf.d/nixl.conf && \
|
||||
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
|
||||
>> /etc/ld.so.conf.d/nixl.conf && \
|
||||
ldconfig
|
||||
|
||||
# Generate RIXL wheel
|
||||
# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
|
||||
# so configure the ROCm wheel variant directly through Meson.
|
||||
# Exclude libcore and libpull from auditwheel: transitive dependencies
|
||||
# that are not shipped in the wheel and vary across base images.
|
||||
RUN cd /opt/rixl && \
|
||||
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
|
||||
contrib/build-wheel.sh && \
|
||||
mkdir -p /app/install && \
|
||||
_ucx_install_dir=${UCX_HOME} \
|
||||
./contrib/build-wheel.sh \
|
||||
--output-dir /app/install \
|
||||
--rocm-dir ${ROCM_PATH} \
|
||||
RUN cd /opt/nixl && \
|
||||
./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
|
||||
CC="ccache gcc" CXX="ccache g++" \
|
||||
uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
|
||||
--python ${PYTHON_VERSION} \
|
||||
-Csetup-args=-Ducx_path=${UCX_HOME} \
|
||||
-Csetup-args=-Dwheel_variant=rocm \
|
||||
-Csetup-args=-Dbuild_tests=false \
|
||||
-Csetup-args=-Dbuild_examples=false && \
|
||||
mkdir -p /tmp/nixl_wheels/repaired /app/install && \
|
||||
auditwheel repair \
|
||||
--exclude 'libamdhip64*' \
|
||||
--exclude 'libcore*' \
|
||||
--exclude 'libpull*' \
|
||||
/tmp/nixl_wheels/nixl_rocm*.whl \
|
||||
--plat manylinux_2_34_$(uname -m) \
|
||||
--wheel-dir /tmp/nixl_wheels/repaired && \
|
||||
./contrib/wheel_add_ucx_plugins.py \
|
||||
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
|
||||
--nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
|
||||
--nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
|
||||
/tmp/nixl_wheels/repaired/*.whl && \
|
||||
cp /tmp/nixl_wheels/repaired/*.whl /app/install
|
||||
|
||||
# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
|
||||
# invalidate the slow ROCShmem build.
|
||||
@@ -660,10 +680,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
|
||||
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
|
||||
fi
|
||||
|
||||
# Install RIXL + DeepEP wheels.
|
||||
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
|
||||
# Install NIXL + DeepEP wheels.
|
||||
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
|
||||
--mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
|
||||
uv pip install --system /rixl_install/*.whl /deep_install/*.whl
|
||||
uv pip install --system /nixl_install/*.whl /deep_install/*.whl
|
||||
|
||||
# Copy ROCShmem runtime libraries.
|
||||
COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
|
||||
@@ -724,6 +744,7 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
|
||||
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
|
||||
# See: https://github.com/ROCm/rocm-libraries/issues/6266
|
||||
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
|
||||
ENV UCX_RMA_PPLN_ENABLE=y
|
||||
|
||||
# ROCm profiler limits workaround.
|
||||
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
|
||||
@@ -796,9 +817,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
|
||||
&& pip uninstall -y vllm \
|
||||
&& uv pip install --system *.whl
|
||||
|
||||
# Install RIXL wheel
|
||||
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
|
||||
uv pip install --system /rixl_install/*.whl
|
||||
# Install NIXL ROCm wheel
|
||||
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
|
||||
uv pip install --system /nixl_install/*.whl
|
||||
|
||||
ARG COMMON_WORKDIR
|
||||
ARG BASE_IMAGE
|
||||
@@ -813,6 +834,7 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
|
||||
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
|
||||
# See: https://github.com/ROCm/rocm-libraries/issues/6266
|
||||
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
|
||||
ENV UCX_RMA_PPLN_ENABLE=y
|
||||
|
||||
ENV TOKENIZERS_PARALLELISM=false
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
|
||||
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
|
||||
ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="v0.1.16.post3"
|
||||
ARG AITER_BRANCH="v0.1.16.post5"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="v1.1.0"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
|
||||
+25
-1
@@ -86,6 +86,29 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
mkdir -p /tmp/hf-xet/dist && \
|
||||
cp dist/*.whl /tmp/hf-xet/dist/
|
||||
|
||||
# Build LLVM 20 from source for llvmlite (system repos ship LLVM 21 which
|
||||
# llvmlite v0.47 does not support; only SystemZ target is needed).
|
||||
FROM base AS llvm20-build
|
||||
ARG LLVM_VERSION=20.1.8
|
||||
WORKDIR /tmp
|
||||
RUN microdnf install -y ninja-build gcc gcc-c++ python3 xz && \
|
||||
curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-${LLVM_VERSION}/llvm-project-${LLVM_VERSION}.src.tar.xz && \
|
||||
tar -xf llvm-project-${LLVM_VERSION}.src.tar.xz && \
|
||||
cmake -G Ninja -S llvm-project-${LLVM_VERSION}.src/llvm -B build \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DCMAKE_INSTALL_PREFIX=/opt/llvm20 \
|
||||
-DLLVM_TARGETS_TO_BUILD="SystemZ" \
|
||||
-DLLVM_ENABLE_RTTI=ON \
|
||||
-DLLVM_BUILD_TOOLS=OFF \
|
||||
-DLLVM_BUILD_UTILS=ON \
|
||||
-DLLVM_BUILD_EXAMPLES=OFF \
|
||||
-DLLVM_BUILD_TESTS=OFF \
|
||||
-DLLVM_INCLUDE_TESTS=OFF \
|
||||
-DLLVM_INCLUDE_EXAMPLES=OFF \
|
||||
-DLLVM_INCLUDE_BENCHMARKS=OFF && \
|
||||
ninja -C build install && \
|
||||
rm -rf build llvm-project-${LLVM_VERSION}.src*
|
||||
|
||||
# Build numba
|
||||
FROM python-install AS numba-builder
|
||||
|
||||
@@ -96,11 +119,13 @@ WORKDIR /tmp
|
||||
|
||||
# Clone all required dependencies
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,from=llvm20-build,source=/opt/llvm20,target=/opt/llvm20 \
|
||||
microdnf install ninja-build gcc gcc-c++ -y && \
|
||||
git clone --recursive https://github.com/numba/llvmlite.git -b v0.47.0 && \
|
||||
git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
|
||||
cd llvmlite && \
|
||||
uv pip install 'cmake<4' 'setuptools<70' numpy && \
|
||||
CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \
|
||||
python setup.py bdist_wheel && \
|
||||
cd ../numba && \
|
||||
if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \
|
||||
@@ -158,7 +183,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
|
||||
OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
|
||||
uv pip install -v \
|
||||
$ARROW_WHL_FILE \
|
||||
$VISION_WHL_FILE \
|
||||
$HF_XET_WHL_FILE \
|
||||
$LLVM_WHL_FILE \
|
||||
|
||||
+13
-13
@@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" {
|
||||
}
|
||||
|
||||
# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
|
||||
# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in
|
||||
# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in
|
||||
# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
|
||||
# stage and is irrelevant when building --target ci_base itself.
|
||||
variable "CI_BASE_IMAGE" {
|
||||
@@ -75,7 +75,7 @@ variable "CI_MAX_JOBS" {
|
||||
# Upstream dependency commit pins -- extracted from Dockerfile.rocm by
|
||||
# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
|
||||
# functions produce no entries when the variable is empty.
|
||||
variable "RIXL_BRANCH" {
|
||||
variable "NIXL_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
variable "RIXL_CACHE_KEY" {
|
||||
variable "NIXL_CACHE_KEY" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" {
|
||||
])
|
||||
}
|
||||
|
||||
# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
|
||||
# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP).
|
||||
# These stages are pinned to specific upstream commit hashes, so cache keys use
|
||||
# those hashes rather than the Buildkite commit. This means the cache persists
|
||||
# across all vLLM commits as long as the upstream dependency pins don't change.
|
||||
@@ -244,16 +244,16 @@ function "get_cache_to_rocm_rust" {
|
||||
function "get_cache_from_rocm_deps" {
|
||||
params = []
|
||||
result = compact([
|
||||
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
|
||||
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
|
||||
ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""),
|
||||
DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""),
|
||||
])
|
||||
}
|
||||
|
||||
function "get_cache_to_rocm_rixl" {
|
||||
function "get_cache_to_rocm_nixl" {
|
||||
params = []
|
||||
result = compact([
|
||||
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
|
||||
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
|
||||
])
|
||||
}
|
||||
|
||||
@@ -372,11 +372,11 @@ variable "CI_BASE_IMAGE_TAG_STABLE" {
|
||||
# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
|
||||
# (e.g., requirements change), these stages are cache hits if their upstream
|
||||
# pins haven't changed -- saving ~35min of compilation.
|
||||
target "rixl-rocm-ci" {
|
||||
target "nixl-rocm-ci" {
|
||||
inherits = ["_common-rocm", "_ci-rocm"]
|
||||
target = "build_rixl"
|
||||
target = "build_nixl"
|
||||
cache-from = get_cache_from_rocm_deps()
|
||||
cache-to = get_cache_to_rocm_rixl()
|
||||
cache-to = get_cache_to_rocm_nixl()
|
||||
output = ["type=cacheonly"]
|
||||
}
|
||||
|
||||
@@ -396,7 +396,7 @@ target "deepep-rocm-ci" {
|
||||
output = ["type=cacheonly"]
|
||||
}
|
||||
|
||||
# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
|
||||
# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
|
||||
# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
|
||||
# files drifts from the remote image label. Per-PR builds then pull the result
|
||||
# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
|
||||
@@ -412,7 +412,7 @@ target "ci-base-rocm-ci" {
|
||||
CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
|
||||
CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
|
||||
]),
|
||||
# Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
|
||||
# Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages
|
||||
# are cache hits even when ci_base itself needs rebuilding.
|
||||
get_cache_from_rocm_deps(),
|
||||
)
|
||||
@@ -424,5 +424,5 @@ target "ci-base-rocm-ci" {
|
||||
# Group for ci_base builds -- exports dependency stage caches alongside the
|
||||
# ci_base image so future rebuilds can reuse them independently.
|
||||
group "ci-base-rocm-ci-with-deps" {
|
||||
targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
|
||||
targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
|
||||
}
|
||||
|
||||
@@ -53,7 +53,7 @@ variable "CI_BASE_IMAGE" {
|
||||
# Upstream dependency commit pins. Plain local bake builds use the Dockerfile
|
||||
# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
|
||||
# overrides) and writes a small HCL override before invoking CI targets.
|
||||
variable "RIXL_BRANCH" {
|
||||
variable "NIXL_BRANCH" {
|
||||
default = ""
|
||||
}
|
||||
|
||||
@@ -106,7 +106,7 @@ target "test-rocm" {
|
||||
output = ["type=docker"]
|
||||
}
|
||||
|
||||
# CI base image target - builds only the ci_base stage (RIXL, DeepEP,
|
||||
# CI base image target - builds only the ci_base stage (NIXL, DeepEP,
|
||||
# torchcodec, requirements, etc.). Used by the weekly scheduled build and
|
||||
# the auto-rebuild trigger when requirements change in a PR.
|
||||
target "ci-base-rocm" {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
"_comment": "Auto-generated from Dockerfile ARGs. Do not edit manually. Run: python tools/generate_versions_json.py",
|
||||
"variable": {
|
||||
"CUDA_VERSION": {
|
||||
"default": "13.0.2"
|
||||
"default": "13.0.3"
|
||||
},
|
||||
"PYTHON_VERSION": {
|
||||
"default": "3.12"
|
||||
@@ -11,10 +11,10 @@
|
||||
"default": "22.04"
|
||||
},
|
||||
"BUILD_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.2-devel-ubuntu22.04"
|
||||
"default": "nvidia/cuda:13.0.3-devel-ubuntu22.04"
|
||||
},
|
||||
"FINAL_BASE_IMAGE": {
|
||||
"default": "nvidia/cuda:13.0.2-base-ubuntu22.04"
|
||||
"default": "nvidia/cuda:13.0.3-base-ubuntu22.04"
|
||||
},
|
||||
"BUILD_OS": {
|
||||
"default": "ubuntu"
|
||||
@@ -68,7 +68,7 @@
|
||||
"default": "true"
|
||||
},
|
||||
"FLASHINFER_VERSION": {
|
||||
"default": "0.6.14"
|
||||
"default": "0.6.15.post1"
|
||||
},
|
||||
"GDRCOPY_CUDA_VERSION": {
|
||||
"default": "12.8"
|
||||
|
||||
@@ -5,7 +5,7 @@ vLLM uses the following environment variables to configure the system:
|
||||
!!! warning
|
||||
Please note that `VLLM_PORT` and `VLLM_HOST_IP` set the port and ip for vLLM's **internal usage**. It is not the port and ip for the API server. If you use `--host $VLLM_HOST_IP` and `--port $VLLM_PORT` to start the API server, it will not work.
|
||||
|
||||
All environment variables used by vLLM are prefixed with `VLLM_`. **Special care should be taken for Kubernetes users**: please do not name the service as `vllm`, otherwise environment variables set by Kubernetes might conflict with vLLM's environment variables, because [Kubernetes sets environment variables for each service with the capitalized service name as the prefix](https://kubernetes.io/docs/concepts/services-networking/service/#environment-variables).
|
||||
Most vLLM-specific environment variables are prefixed with `VLLM_` (a handful of standard names — for example `CUDA_VISIBLE_DEVICES`, `MAX_JOBS`, `S3_ACCESS_KEY_ID`/`S3_SECRET_ACCESS_KEY`/`S3_ENDPOINT_URL`, `DO_NOT_TRACK`, `NO_COLOR` — are also read directly when set). **Special care should be taken for Kubernetes users**: please do not name the service as `vllm`, otherwise environment variables set by Kubernetes might conflict with vLLM's environment variables, because [Kubernetes sets environment variables for each service with the capitalized service name as the prefix](https://kubernetes.io/docs/concepts/services-networking/service/#environment-variables).
|
||||
|
||||
```python
|
||||
--8<-- "vllm/envs.py:env-vars-definition"
|
||||
|
||||
@@ -6,7 +6,11 @@ vLLM maintains a per-commit wheel repository (commonly referred to as "nightly")
|
||||
|
||||
### Wheel Building
|
||||
|
||||
Wheels are built in the `Release` pipeline (`.buildkite/release-pipeline.yaml`) after a PR is merged into the main branch, with multiple variants:
|
||||
Wheels are built in the `Release` pipeline
|
||||
(`.buildkite/release-pipeline.yaml`) after a PR is merged into the main branch.
|
||||
Regular builds produce the CUDA 13.0 wheels for x86_64 and aarch64. Additional
|
||||
wheel variants and ROCm builds can be unblocked on demand and run automatically
|
||||
when `NIGHTLY=1`:
|
||||
|
||||
- **Backend variants**: `cpu` and `cuXXX` (e.g., `cu129`, `cu130`).
|
||||
- **Architecture variants**: `x86_64` and `aarch64`.
|
||||
|
||||
@@ -164,8 +164,8 @@ Priority is **1 = highest** (tried first).
|
||||
| `FLASHINFER` | XQA† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ❌ | ❌ | ❌ | ✅ | Decoder | 9.0 |
|
||||
| `FLASHINFER` | trtllm-gen† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `nvfp4` | 16, 32, 64, 128, 256, 512, 1024 | 64, 128, 256, 512 | ✅ | ✅ | ❌ | ✅ | Decoder | 10.x |
|
||||
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ❌ | ✅ | All | ≥8.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | 9.x |
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
|
||||
@@ -205,7 +205,7 @@ hardware and configuration.
|
||||
|
||||
| Backend | Description | Dtypes | Compute Cap. | Notes |
|
||||
| ------- | ----------- | ------ | ------------ | ----- |
|
||||
| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) |
|
||||
| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) |
|
||||
| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only |
|
||||
| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
|
||||
| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
|
||||
|
||||
@@ -306,7 +306,7 @@ Supported quantization scheme/hardware combinations:
|
||||
|
||||
- Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py)
|
||||
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
|
||||
- CUDA/HIP kernels: [`csrc/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_quant_kernels.cu)
|
||||
- CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu)
|
||||
|
||||
### SiLU+Mul + Quantization (`fuse_act_quant`)
|
||||
|
||||
@@ -332,7 +332,7 @@ Supported quantization scheme/hardware combinations:
|
||||
- Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py)
|
||||
- ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
|
||||
- CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/)
|
||||
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
|
||||
- Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
|
||||
|
||||
### RMSNorm + Padding (`fuse_act_padding`)
|
||||
|
||||
|
||||
@@ -107,6 +107,7 @@ Batch invariance has been tested and verified on the following models:
|
||||
- **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-3B-Instruct` for example
|
||||
- **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b`
|
||||
- **Mistral**: `mistralai/Mistral-7B-v0.3`
|
||||
- **Phi series**: `microsoft/Phi-3.5-mini-instruct`
|
||||
|
||||
Other models may also work, but these have been explicitly validated. If you encounter issues with a specific model, please report them on the [GitHub issue tracker](https://github.com/vllm-project/vllm/issues/new/choose).
|
||||
|
||||
|
||||
@@ -68,13 +68,14 @@ vllm serve <model> \
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `spec_name` | no | `CPUOffloadingSpec` | both | Set to `TieringOffloadingSpec` for multi-tier. |
|
||||
| `cpu_bytes_to_use` | yes | — | both | Total bytes of host memory reserved for the CPU tier across all workers (not per-worker). |
|
||||
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. |
|
||||
| `block_size` | no | GPU block size | both | Offloaded block size in tokens; must be a multiple of the GPU block size. Mutually exclusive with `blocks_per_chunk`. |
|
||||
| `blocks_per_chunk` | no | `1` | both | Offloaded chunk size in GPU blocks; must be > 0. Alternative to `block_size` for models whose KV cache groups have different block sizes. |
|
||||
| `eviction_policy` | no | `lru` | both | Primary tier policy: `lru` or `arc`. |
|
||||
| `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. |
|
||||
| `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. |
|
||||
| `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). |
|
||||
| `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. |
|
||||
| `self_describing_kv_events` | no | `false` | single-tier | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. Currently rejected by `TieringOffloadingSpec`. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
|
||||
| `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Full-attention groups only; sliding-window/SSM groups keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. |
|
||||
| `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). |
|
||||
|
||||
## Secondary Tiers
|
||||
@@ -83,9 +84,11 @@ Each entry in `secondary_tiers` is a dict with a required `type` field plus tier
|
||||
|
||||
The filesystem and object-store tiers can publish hash-only `BlockStored` KV events for blocks they successfully store, tagged with a stable per-tier `medium` (`FS` for the filesystem tier, `OBJ` for the object-store tier). Set `enable_kv_events: true` in the tier's entry to opt in; events are published only when KV cache events are also enabled globally via `--kv-events-config`.
|
||||
|
||||
Set the optional `locality` tier field to `LOCAL` or `REMOTE` to describe the tier's storage location relative to the publishing vLLM instance. `LOCAL` marks storage local to that instance, while `REMOTE` marks storage that is not local to it. When the setting is omitted, locality is unspecified. vLLM does not infer it from the tier type, so an OBJ tier is not implicitly `REMOTE`. A KV event includes `locality` only when the tier explicitly configures it. This metadata describes the tier property without implying that a consumer can already route requests to its blocks.
|
||||
|
||||
### Filesystem (FS)
|
||||
|
||||
The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage.
|
||||
The filesystem tier (`type: "fs"`) writes blocks to a filesystem directory.
|
||||
|
||||
| Key | Required | Default | Notes |
|
||||
| --- | --- | --- | --- |
|
||||
@@ -94,6 +97,7 @@ The filesystem tier (`type: "fs"`) writes blocks to a directory on local storage
|
||||
| `n_read_threads` | no | `16` | Read-priority I/O threads (load path). |
|
||||
| `n_write_threads` | no | `16` | Write-priority I/O threads (store path). |
|
||||
| `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `FS`) for successfully stored blocks. Requires KV cache events to be enabled globally. |
|
||||
| `locality` | no | unspecified | `LOCAL` or `REMOTE` relative to the publishing vLLM instance. Included in the tier's KV events only when explicitly configured. |
|
||||
|
||||
Each thread group prefers its own queue but pulls from the other when its primary queue is empty, so a write-heavy or read-heavy burst won't leave the off-priority queue waiting. Size the totals to your storage's effective concurrency.
|
||||
|
||||
@@ -134,6 +138,7 @@ The object-store tier (`type: "obj"`) offloads blocks to an S3-compatible object
|
||||
| `prefix` | no | `""` | Key prefix prepended to all object keys. |
|
||||
| `io_threads` | no | `4` | Number of NIXL OBJ backend I/O threads. |
|
||||
| `enable_kv_events` | no | `false` | Publish `BlockStored` KV events (medium `OBJ`) for successfully stored blocks. Requires KV cache events to be enabled globally. |
|
||||
| `locality` | no | unspecified | `LOCAL` or `REMOTE` relative to the publishing vLLM instance. Included in the tier's KV events only when explicitly configured; OBJ does not imply `REMOTE`. |
|
||||
|
||||
`store_config` fields:
|
||||
|
||||
@@ -175,7 +180,7 @@ Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them on
|
||||
|
||||
- `cpu_bytes_to_use`: a bigger CPU tier means fewer trips to slower secondary tiers and a higher hit rate. The value is total across all workers, not per-worker. Leave headroom for the rest of the host workload.
|
||||
- For single-tier (CPU-only) setups, set `cpu_bytes_to_use` larger than the aggregate GPU KV cache. Because offloading is immediate, a smaller CPU tier just mirrors what the GPU already holds and adds no hit rate.
|
||||
- `block_size`: larger offloaded blocks reduce per-block bookkeeping overhead but increase the granularity of lookups. Must be a multiple of the GPU block size.
|
||||
- `block_size` / `blocks_per_chunk`: larger offloaded chunks reduce per-block bookkeeping overhead but increase the granularity of lookups.
|
||||
- FS thread counts: tune `n_read_threads` and `n_write_threads` to the parallelism your storage can sustain. Reads are latency-sensitive on the prefill path, so prefer more read threads when prefill hit rates are high.
|
||||
- Sharing `root_dir` across runs: runs with the same model, `block_size`, parallelism layout, and dtype share files under the same `<digest>` subdirectory. Changing any of these produces a new subdirectory; old ones are orphaned but harmless. Delete them to reclaim disk.
|
||||
|
||||
|
||||
@@ -13,11 +13,7 @@ Install the NIXL library: `uv pip install nixl`, as a quick start on Nvidia plat
|
||||
- Refer to [NIXL official repository](https://github.com/ai-dynamo/nixl) for more installation instructions
|
||||
- The specified required NIXL version can be found in [requirements/kv_connectors.txt](../../requirements/kv_connectors.txt) and other relevant config files
|
||||
|
||||
For ROCm platform, the [ROCm docker file](../../docker/Dockerfile.rocm) includes RIXL and ucx already.
|
||||
|
||||
- Refer to [RIXL official repository](https://github.com/rocm/rixl) for more information
|
||||
- The supportive libraries for RIXL can be found in [requirements/kv_connectors_rocm.txt](../../requirements/kv_connectors_rocm.txt)
|
||||
- In the future we may remove RIXL from docker image file and users will be able to install from pre-compiled binary packages
|
||||
For ROCm, the [ROCm Dockerfile](../../docker/Dockerfile.rocm) builds NIXL and UCX with ROCm support from source.
|
||||
|
||||
For non-cuda platform, please install nixl with ucx build from source, instructed as below.
|
||||
|
||||
|
||||
@@ -315,7 +315,7 @@ vLLM CPU supports data parallel (DP), tensor parallel (TP) and pipeline parallel
|
||||
- vLLM CPU supports quantizations:
|
||||
- AWQ (x86 only)
|
||||
- GPTQ (x86 only)
|
||||
- compressed-tensor INT8 W8A8 (x86, s390x)
|
||||
- compressed-tensor INT8 W8A8 (x86 only)
|
||||
|
||||
### Why do I see `get_mempolicy: Operation not permitted` when running in Docker?
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ Currently, the CPU implementation for s390x architecture supports FP32, BF16 and
|
||||
- OS: `Linux`
|
||||
- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools
|
||||
- Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above.
|
||||
- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless`
|
||||
- Build from source python packages (no pre-built s390x wheels): `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`
|
||||
|
||||
--8<-- [end:requirements]
|
||||
--8<-- [start:set-up-using-python]
|
||||
@@ -28,13 +28,24 @@ Install the following packages from the package manager before building the vLLM
|
||||
|
||||
```bash
|
||||
dnf install -y \
|
||||
which procps findutils tar vim git gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
|
||||
which procps findutils tar vim git patch xz ninja-build \
|
||||
gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
|
||||
libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \
|
||||
openssl-devel openblas openblas-devel autoconf automake libtool cmake numpy libsndfile \
|
||||
clang llvm-devel llvm-static clang-devel
|
||||
```
|
||||
|
||||
Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation.
|
||||
Build and install `numactl` from source:
|
||||
|
||||
```bash
|
||||
curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz
|
||||
tar -xvzf v2.0.19.tar.gz
|
||||
cd numactl-2.0.19
|
||||
./autogen.sh && ./configure && make && make install
|
||||
cd ..
|
||||
```
|
||||
|
||||
Install rust>=1.80 which is needed for `outlines-core`, `uvloop`, and `hf-xet` python packages installation.
|
||||
|
||||
```bash
|
||||
curl https://sh.rustup.rs -sSf | sh -s -- -y && \
|
||||
@@ -44,26 +55,79 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \
|
||||
Execute the following commands to build and install vLLM from source.
|
||||
|
||||
!!! tip
|
||||
Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM.
|
||||
Pre-built wheels are not available for s390x for the following packages. Build them from source before building vLLM: `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`.
|
||||
See `docker/Dockerfile.s390x` for exact versions and build commands used in each multi-stage build.
|
||||
|
||||
!!! note "LLVM 20 required for llvmlite"
|
||||
`llvmlite v0.47` requires LLVM 20, but UBI 9.6 repos ship LLVM 21 which is
|
||||
not compatible. You must build LLVM 20 from source before building `llvmlite`:
|
||||
|
||||
```bash
|
||||
curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-20.1.8/llvm-project-20.1.8.src.tar.xz
|
||||
tar -xf llvm-project-20.1.8.src.tar.xz
|
||||
cmake -G Ninja -S llvm-project-20.1.8.src/llvm -B llvm-build \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DCMAKE_INSTALL_PREFIX=/opt/llvm20 \
|
||||
-DLLVM_TARGETS_TO_BUILD="SystemZ" \
|
||||
-DLLVM_ENABLE_RTTI=ON \
|
||||
-DLLVM_BUILD_TOOLS=OFF \
|
||||
-DLLVM_BUILD_UTILS=ON \
|
||||
-DLLVM_BUILD_EXAMPLES=OFF \
|
||||
-DLLVM_BUILD_TESTS=OFF \
|
||||
-DLLVM_INCLUDE_TESTS=OFF \
|
||||
-DLLVM_INCLUDE_EXAMPLES=OFF \
|
||||
-DLLVM_INCLUDE_BENCHMARKS=OFF
|
||||
ninja -C llvm-build install
|
||||
```
|
||||
|
||||
Then build `llvmlite` pointing to LLVM 20:
|
||||
|
||||
```bash
|
||||
CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \
|
||||
python setup.py bdist_wheel
|
||||
```
|
||||
|
||||
```bash
|
||||
uv pip install -v \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt \
|
||||
--torch-backend cpu \
|
||||
--index-strategy unsafe-best-match && \
|
||||
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
|
||||
uv pip install dist/*.whl
|
||||
uv pip install -v \
|
||||
/path/to/torchvision.whl \
|
||||
/path/to/llvmlite.whl \
|
||||
/path/to/numba.whl \
|
||||
/path/to/opencv_python_headless.whl \
|
||||
/path/to/hf_xet.whl \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt \
|
||||
--torch-backend cpu \
|
||||
--index-strategy unsafe-best-match && \
|
||||
VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
|
||||
uv pip install dist/*.whl
|
||||
```
|
||||
|
||||
??? console "pip"
|
||||
```bash
|
||||
pip install -v \
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt \
|
||||
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
|
||||
pip install dist/*.whl
|
||||
pip install -v \
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu \
|
||||
/path/to/torchvision.whl \
|
||||
/path/to/llvmlite.whl \
|
||||
/path/to/numba.whl \
|
||||
/path/to/opencv_python_headless.whl \
|
||||
/path/to/hf_xet.whl \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt && \
|
||||
VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
|
||||
pip install dist/*.whl
|
||||
```
|
||||
|
||||
!!! warning "Protobuf workaround for s390x"
|
||||
The C++ protobuf extension crashes on s390x. After installation, set the
|
||||
following environment variable and remove the C++ extensions:
|
||||
|
||||
```bash
|
||||
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
|
||||
|
||||
# Remove C++ protobuf extensions that crash on s390x
|
||||
SITE_PKGS=$(python -c "import site; print(site.getsitepackages()[0])")
|
||||
rm -rf "$SITE_PKGS/google/_upb/"*.so \
|
||||
"$SITE_PKGS/google/protobuf/pyext/"*.so 2>/dev/null || true
|
||||
```
|
||||
|
||||
--8<-- [end:build-wheel-from-source]
|
||||
@@ -80,19 +144,20 @@ docker build -f docker/Dockerfile.s390x \
|
||||
|
||||
# Launch OpenAI server
|
||||
docker run --rm \
|
||||
--privileged true \
|
||||
--security-opt seccomp=unconfined \
|
||||
--cap-add SYS_NICE \
|
||||
--shm-size 4g \
|
||||
-p 8000:8000 \
|
||||
-e VLLM_CPU_KVCACHE_SPACE=<KV cache space> \
|
||||
-e VLLM_CPU_OMP_THREADS_BIND=<CPU cores for inference> \
|
||||
vllm-cpu-env \
|
||||
--model meta-llama/Llama-3.2-1B-Instruct \
|
||||
--dtype float \
|
||||
--dtype bfloat16 \
|
||||
other vLLM OpenAI server arguments
|
||||
```
|
||||
|
||||
!!! tip
|
||||
An alternative of `--privileged true` is `--cap-add SYS_NICE --security-opt seccomp=unconfined`.
|
||||
Alternatively, `--privileged=true` also works but is broader and not generally recommended.
|
||||
|
||||
--8<-- [end:build-image-from-source]
|
||||
--8<-- [start:extra-information]
|
||||
|
||||
@@ -27,7 +27,7 @@ Currently, there are no pre-built XPU wheels.
|
||||
|
||||
- First, install required [driver](https://dgpu-docs.intel.com/driver/installation.html#installing-gpu-drivers).
|
||||
- Second, install Python packages for vLLM XPU backend building (Intel OneAPI dependencies are installed automatically as part of `torch-xpu`, see [PyTorch XPU get started](https://docs.pytorch.org/docs/stable/notes/get_start_xpu.html)):
|
||||
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.14.37833.4) release, to avoid potential compatibility issue.
|
||||
- Start from vllm-xpu-kernels v0.1.10, we recommend user upgrade driver to [compute runtime 26.18](https://github.com/intel/compute-runtime/releases/tag/26.18.38308.1) release, to avoid potential compatibility issue.
|
||||
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
@@ -58,7 +58,40 @@ VLLM_TARGET_DEVICE=xpu pip install --no-build-isolation -e . -v
|
||||
--8<-- [end:build-wheel-from-source]
|
||||
--8<-- [start:pre-built-images]
|
||||
|
||||
Currently, we release prebuilt XPU images at docker [hub](https://hub.docker.com/r/intel/vllm/tags) based on vLLM released version. For more information, please refer release [note](https://github.com/intel/ai-containers/blob/main/vllm).
|
||||
vLLM offers official Docker images for deployment.
|
||||
The images can be used to run OpenAI compatible server and are available on Docker Hub as [vllm/vllm-openai-xpu](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
|
||||
|
||||
- `vllm/vllm-openai-xpu:latest` — stable release, available starting from v0.26.0
|
||||
- `vllm/vllm-openai-xpu:nightly` — preview build from the latest development branch, use this if you want the latest features and fixes
|
||||
|
||||
```bash
|
||||
docker run --rm \
|
||||
--network=host \
|
||||
--device /dev/dri:/dev/dri \
|
||||
-v /dev/dri/by-path:/dev/dri/by-path \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
--ipc=host \
|
||||
--privileged \
|
||||
vllm/vllm-openai-xpu:<tag> \
|
||||
--model Qwen/Qwen3-0.6B
|
||||
```
|
||||
|
||||
To use the docker image as base for development, you can launch it in interactive session through overriding the entrypoint.
|
||||
|
||||
???+ console "Commands"
|
||||
```bash
|
||||
docker run --rm -it \
|
||||
--network=host \
|
||||
--device /dev/dri:/dev/dri \
|
||||
-v /dev/dri/by-path:/dev/dri/by-path \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
--ipc=host \
|
||||
--privileged \
|
||||
--entrypoint /bin/bash \
|
||||
vllm/vllm-openai-xpu:<tag>
|
||||
```
|
||||
|
||||
--8<-- [end:pre-built-images]
|
||||
--8<-- [start:build-image-from-source]
|
||||
|
||||
@@ -65,6 +65,15 @@ This guide will help you quickly get started with vLLM to perform:
|
||||
!!! tip
|
||||
A nightly Docker image is also available as [vllm/vllm-openai-rocm:nightly](https://hub.docker.com/r/vllm/vllm-openai-rocm/tags) for testing the latest development builds.
|
||||
|
||||
=== "Intel GPU"
|
||||
|
||||
vLLM supports Intel GPUs through the XPU backend. Pre-built XPU wheels will be available soon.
|
||||
|
||||
Official Docker images for Intel GPUs are added to the vLLM release starting from v0.26.0. Nightly Docker image is also available as [vllm/vllm-openai-xpu:nightly](https://hub.docker.com/r/vllm/vllm-openai-xpu/tags).
|
||||
|
||||
!!! tip
|
||||
For more detailed instructions, including building from source and Docker image setup, please refer to the [GPU installation guide](installation/gpu.md) and select the "Intel XPU" tab.
|
||||
|
||||
=== "Google TPU"
|
||||
|
||||
To run vLLM on Google TPUs, you need to install the `vllm-tpu` package.
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
// Reo.Dev documentation tracking
|
||||
// https://docs.reo.dev/integrations/input-sources/developer-insights/documentation
|
||||
!function(){var e,t,n;e="d5c4337961ef0ac",t=function(){Reo.init({clientID:"d5c4337961ef0ac", enableThirdPartyTracking: true})},(n=document.createElement("script")).src="https://static.reo.dev/"+e+"/reo.js",n.defer=!0,n.onload=t,document.head.appendChild(n)}();
|
||||
@@ -1,5 +1,5 @@
|
||||
{% extends "base.html" %}
|
||||
|
||||
{% block announce %}
|
||||
<p>You are viewing the latest developer preview docs. <a href="https://docs.vllm.ai/en/stable/">Click here</a> to view docs for the latest stable release.</p>
|
||||
<p>You are viewing the latest developer preview docs. <a href="https://docs.vllm.ai/en/stable/{{ page.url }}">Click here</a> to view docs for the latest stable release.</p>
|
||||
{% endblock %}
|
||||
|
||||
@@ -31,10 +31,8 @@
|
||||
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
|
||||
| chuhac/TeleChat2-35B | LlamaForCausalLM (TeleChat2 based on Llama arch) | ✅ | | |
|
||||
| 01-ai/Yi1.5-34B-Chat | YiForCausalLM | ✅ | | |
|
||||
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
|
||||
| deepseek-ai/DeepSeek-Coder-33B-base | DeepSeekCoderForCausalLM | ✅ | | |
|
||||
| meta-llama/Llama-2-13b-chat-hf | LlamaForCausalLM | ✅ | | |
|
||||
| THUDM/CodeGeex4-All-9B | CodeGeexForCausalLM | ✅ | | |
|
||||
| Qwen/Qwen1.5-14B-Chat | QwenForCausalLM | ✅ | | |
|
||||
| Qwen/Qwen1.5-32B-Chat | QwenForCausalLM | ✅ | | |
|
||||
| RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic | LlamaForCausalLM | | ✅ | |
|
||||
|
||||
@@ -220,8 +220,6 @@ For multi-node deployment, add these EPLB flags to each node's command. We recom
|
||||
|
||||
- Use simulator flags `VLLM_MOE_ROUTING_SIMULATION_STRATEGY=uniform_random` and `VLLM_RANDOMIZE_DP_DUMMY_INPUTS=1` so token routing is balanced across EP ranks.
|
||||
|
||||
- Increasing `VLLM_MOE_DP_CHUNK_SIZE` may increase throughput by increasing the maximum batch size for inter-rank token transfers. This may cause DeepEP to throw `assert self.nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2`, which can be fixed by increasing environment variable `NVSHMEM_QP_DEPTH`.
|
||||
|
||||
## Disaggregated Serving (Prefill/Decode Split)
|
||||
|
||||
For production deployments requiring strict SLA guarantees for time-to-first-token and inter-token latency, disaggregated serving allows independent scaling of prefill and decode operations.
|
||||
|
||||
@@ -137,8 +137,12 @@ For further details on renderer APIs, please refer to [this page](renderer.md).
|
||||
|
||||
### Derenderer APIs
|
||||
|
||||
- `/v1/completions/derender` - Derenderer completion requests
|
||||
- `/v1/chat/completions/derender` - Derenderer chat completion requests
|
||||
For further details on derenderer APIs, please refer to [this page](derenderer.md).
|
||||
|
||||
- [Chat Completions Derender API](derenderer.md) (`/v1/chat/completions/derender`)
|
||||
- Derender chat completion requests
|
||||
- [Completions Derender API](derenderer.md) (`/v1/completions/derender`)
|
||||
- Derender completion requests
|
||||
|
||||
## Tokenize APIs
|
||||
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
# Derenderer APIs
|
||||
|
||||
The derenderer API is the post processing counterpart to the [Renderer APIs](renderer.md). Where `/render` turns a request into token ID (preprocessing), `/derender` turns generated token IDs back into a fully formed OpenAI compatible response (detokenization, reasoning parsing, tool call parsing), all without a GPU.
|
||||
|
||||
This closes the loop for a token-in / token-out engine in disaggregated serving:
|
||||
|
||||
- **GPU less post processing**: Detokenization, reasoning parsing, and tool call parsing run on the same GPU less frontend that hosts `/render`
|
||||
- **Parser parity**: The derenderer reuses vLLM's tool and reasoning parsers, so a disaggregated deployment produces the same `content`/`reasoning`/ `tool_calls` split as a standard `vllm serve` server
|
||||
- **Non-streaming**: The endpoints expect a complete `GenerateResponse` with all token IDs present and perform one-shot parsing. Streaming derender would require a separate endpoint design and is not currently supported but is in the pipeline
|
||||
|
||||
Both endpoints are hosted by the GPU less rendering server started with [`vllm launch render`](../../cli/launch/render.md), alongside the `/render`
|
||||
endpoints.
|
||||
|
||||
## Pipeline
|
||||
|
||||
```text
|
||||
render generate derender
|
||||
request ───────────────▶ token_ids ─────────▶ token_ids ──────────▶ response
|
||||
(chat / (GPU less) (token-in / (GPU less) (OpenAI
|
||||
completion) │ token-out engine) ▲ compatible)
|
||||
└─────────────── request + prompt_tokens ──┘
|
||||
```
|
||||
|
||||
The derender step needs more than the engine's `token_ids`. It also consumes the original `chat_request`/`completion_request` and `prompt_tokens` carried over from the render step (see [Request format](#request-format)) so the tool and reasoning parsers have the context they need.
|
||||
|
||||
## API Reference
|
||||
|
||||
- Chat Completions Derender API (`/v1/chat/completions/derender`)
|
||||
- Post process a single `GenerateResponse` into a `ChatCompletionResponse`
|
||||
- Completions Derender API (`/v1/completions/derender`)
|
||||
- Post process a list of `GenerateResponse` objects (one per prompt) into a `CompletionResponse`
|
||||
|
||||
## Request format
|
||||
|
||||
Each request wraps the engine's `GenerateResponse`(s) together with the caller metadata needed to reconstruct the final response without a GPU.
|
||||
|
||||
`/v1/chat/completions/derender`:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
--8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-chat-request"
|
||||
```
|
||||
|
||||
`/v1/completions/derender`:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
--8<-- "vllm/entrypoints/scale_out/token_in_token_out/protocol.py:derender-completion-request"
|
||||
```
|
||||
|
||||
Oversized payloads are rejected with a `400` before any `tokenizer.decode()` or parser runs.
|
||||
|
||||
## Example
|
||||
|
||||
The example below drives the full `render → generate → derender` round trip for a chat request against a GPU less render server (`/render`, `/derender`) and a token-in / token-out engine (`/inference/v1/generate`).
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
MODEL = "meta-llama/Llama-3.2-1B-Instruct"
|
||||
RENDER = "http://localhost:8100" # vllm launch render ...
|
||||
ENGINE = "http://localhost:8200" # token-in / token-out engine
|
||||
|
||||
chat_request = {
|
||||
"model": MODEL,
|
||||
"messages": [{"role": "user", "content": "What is 2+2?"}],
|
||||
"max_tokens": 32,
|
||||
}
|
||||
|
||||
with httpx.Client(timeout=60.0) as client:
|
||||
# 1. Render: request -> token IDs (GPU less)
|
||||
generate_request = client.post(
|
||||
f"{RENDER}/v1/chat/completions/render", json=chat_request
|
||||
).json()
|
||||
prompt_tokens = len(generate_request["token_ids"])
|
||||
|
||||
# 2. Generate: token IDs -> token IDs (token-in / token-out engine)
|
||||
generate_response = client.post(
|
||||
f"{ENGINE}/inference/v1/generate", json=generate_request
|
||||
).json()
|
||||
|
||||
# 3. Derender: token IDs -> ChatCompletionResponse (GPU less)
|
||||
response = client.post(
|
||||
f"{RENDER}/v1/chat/completions/derender",
|
||||
json={
|
||||
"model": MODEL,
|
||||
"generate_response": generate_response,
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"chat_request": chat_request,
|
||||
},
|
||||
).json()
|
||||
|
||||
print(response["choices"][0]["message"]["content"])
|
||||
```
|
||||
|
||||
Passing `chat_request` lets the derenderer run the configured tool and reasoning parsers. This means `response["choices"][0]["message"]` carries the same `content` / `reasoning` / `tool_calls` split a `vllm serve` server would produce. Omit `chat_request` for plain detokenization only.
|
||||
@@ -12,3 +12,5 @@ Our renderer API is designed to disaggregate the render phase(preprocessing) and
|
||||
- Render completion requests
|
||||
- [Chat Completions Render API](renderer.md) (`/v1/chat/completions/render`)
|
||||
- Render chat completions
|
||||
|
||||
For the post processing counterpart that turns generated token IDs back into OpenAI compatible responses, see the [Derenderer APIs](derenderer.md).
|
||||
|
||||
@@ -67,7 +67,7 @@ The Transcriptions API supports uploading audio files in various formats includi
|
||||
- `response_format`: Format of the response ("json", "text") (optional)
|
||||
- `temperature`: Sampling temperature between 0 and 1 (optional)
|
||||
|
||||
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/protocol.py#L2182).
|
||||
For the complete list of supported parameters including sampling parameters and vLLM extensions, see the [protocol definitions](https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/speech_to_text/transcription/protocol.py).
|
||||
|
||||
**Response Format:**
|
||||
|
||||
|
||||
@@ -155,8 +155,10 @@ When `--api-key` is configured, the following `/v1` endpoints require Bearer tok
|
||||
- `/v1/chat/completions` - Chat completions
|
||||
- `/v1/chat/completions/batch` - Batch chat completions
|
||||
- `/v1/chat/completions/render` - Render chat completion requests
|
||||
- `/v1/chat/completions/derender` - Derender chat completion requests
|
||||
- `/v1/completions` - Text completions
|
||||
- `/v1/completions/render` - Render completion requests
|
||||
- `/v1/completions/derender` - Derender completion requests
|
||||
- `/v1/embeddings` - Generate embeddings
|
||||
- `/v1/audio/transcriptions` - Audio transcription
|
||||
- `/v1/audio/translations` - Audio translation
|
||||
|
||||
@@ -210,8 +210,31 @@ async def stream_decode_response(session, response, request_id):
|
||||
await session.close()
|
||||
|
||||
|
||||
def example_round_robin_dp_loader(request_number, dp_size):
|
||||
return request_nums % dp_size
|
||||
def flat_interleaved_dp_route(request_number, instances):
|
||||
"""Flat round-robin over the full (instance, dp_rank) slot space.
|
||||
|
||||
ONE counter over (n_instances * dp_size) slots, so instance-selection and
|
||||
DP-rank-selection are derived from the SAME index and can never alias. The
|
||||
previous scheme computed instance = req % n and rank = req % dp from the
|
||||
same counter with n | dp, which locked each instance to a stride-n subset
|
||||
of its ranks (e.g. 2 prefill instances -> 4 of 8 ranks each -> half the
|
||||
GPUs never receive a request, so the deployment falsely appears not to
|
||||
scale).
|
||||
|
||||
Interleaved order — inst0_r0, inst1_r0, inst0_r1, inst1_r1, ... — so
|
||||
consecutive requests alternate instances AND every rank gets walked.
|
||||
|
||||
Assumes homogeneous dp_size across a role's instances (true for the
|
||||
DP<->DP and DP<->TP deployments this proxy targets). Returns
|
||||
(instance_index, dp_rank); dp_rank is None when dp_size == 1 (e.g. a TP
|
||||
decode), which avoids forwarding an out-of-range data-parallel rank.
|
||||
"""
|
||||
n = len(instances)
|
||||
dp = instances[0]["dp_size"]
|
||||
slot = (request_number - 1) % (n * dp)
|
||||
inst_idx = slot % n
|
||||
dp_rank = (slot // n) if dp > 1 else None
|
||||
return inst_idx, dp_rank
|
||||
|
||||
|
||||
@app.route("/health", methods=["GET"])
|
||||
@@ -252,18 +275,21 @@ async def handle_request(api: str, request: Request):
|
||||
503,
|
||||
)
|
||||
)
|
||||
pid = request_nums % len(prefill_instances)
|
||||
did = request_nums % len(decode_instances)
|
||||
# Flat interleaved round-robin (see flat_interleaved_dp_route): ONE
|
||||
# counter over the full (instance, dp_rank) slot space per role, so
|
||||
# instance-selection and DP-rank-selection derive from the same index
|
||||
# and can never alias. The old scheme keyed both on request_nums with
|
||||
# n_instances | dp_size, stranding half the ranks (e.g. in 2P_DP8EP).
|
||||
pid, selected_prefill_dp_rank = flat_interleaved_dp_route(
|
||||
request_nums, prefill_instances
|
||||
)
|
||||
# Decode instance selection uses the same interleaved walk; in READ
|
||||
# mode the decode reads KV from selected_prefill_dp_rank, so the
|
||||
# decode's own dp_rank is not forwarded here.
|
||||
did, _ = flat_interleaved_dp_route(request_nums, decode_instances)
|
||||
prefill_instance_endpoint = prefill_instances[pid]
|
||||
decode_instance_endpoint = decode_instances[did]
|
||||
|
||||
selected_prefill_dp_rank = None
|
||||
if prefill_instance_endpoint["dp_size"] > 1:
|
||||
selected_prefill_dp_rank = example_round_robin_dp_loader(
|
||||
request_nums // len(prefill_instance_endpoint),
|
||||
prefill_instance_endpoint["dp_size"],
|
||||
)
|
||||
|
||||
# Embed both zmq_addresses in the request_id so the connector can parse
|
||||
# the peer's host/ports from it, similar to P2P-NCCL
|
||||
uid = str(uuid.uuid4()).replace("-", "")
|
||||
@@ -427,9 +453,33 @@ if __name__ == "__main__":
|
||||
args = parser.parse_args()
|
||||
|
||||
t = start_service_discovery("0.0.0.0", 36367)
|
||||
app.debug = True
|
||||
# High-concurrency hardening. Quart's app.run() uses a shallow listen
|
||||
# backlog (100) and, with app.debug=True, adds per-request overhead that
|
||||
# starves the single accept loop. Under a burst of ~512 simultaneous client
|
||||
# connections the backlog overflows and the kernel RSTs the excess, so
|
||||
# clients see "ClientOSError: [Errno 104] Connection reset by peer" before
|
||||
# any response (~16% request loss at c=512). Serve via hypercorn with debug
|
||||
# OFF and a deep backlog so the burst QUEUES (higher TTFT) instead of being
|
||||
# reset -> 100% request success.
|
||||
app.debug = False
|
||||
app.config["BODY_TIMEOUT"] = 360000
|
||||
app.config["RESPONSE_TIMEOUT"] = 360000
|
||||
|
||||
app.run(host="0.0.0.0", port=args.port)
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
from hypercorn.asyncio import serve as _hypercorn_serve
|
||||
from hypercorn.config import Config as _HypercornConfig
|
||||
|
||||
_hcfg = _HypercornConfig()
|
||||
_hcfg.bind = [f"0.0.0.0:{args.port}"]
|
||||
# Deep listen backlog so a wide connection burst queues, not RSTs. NOTE:
|
||||
# effective backlog is capped by the host's net.core.somaxconn (proxy runs
|
||||
# --network host); kernel 6.x defaults to 4096. Override via
|
||||
# PROXY_LISTEN_BACKLOG.
|
||||
_hcfg.backlog = int(os.environ.get("PROXY_LISTEN_BACKLOG", "4096"))
|
||||
# Long-lived SSE streams (8k1k decode ~5 min): never reap on keepalive.
|
||||
_hcfg.keep_alive_timeout = 360000.0
|
||||
|
||||
asyncio.run(_hypercorn_serve(app, _hcfg))
|
||||
t.join()
|
||||
|
||||
@@ -42,12 +42,16 @@ class BlockStored(KVCacheEvent):
|
||||
"""
|
||||
|
||||
group_idx: int | None = None
|
||||
kv_cache_spec_kind: str | None = None
|
||||
kv_cache_spec_sliding_window: int | None = None
|
||||
locality: str | None = None
|
||||
|
||||
|
||||
class BlockRemoved(KVCacheEvent):
|
||||
block_hashes: list[ExternalBlockHash]
|
||||
medium: str | None
|
||||
group_idx: int | None = None
|
||||
locality: str | None = None
|
||||
|
||||
|
||||
class AllBlocksCleared(KVCacheEvent):
|
||||
|
||||
@@ -17,6 +17,7 @@ from transformers import AutoProcessor, AutoTokenizer
|
||||
from vllm import LLM, EngineArgs, SamplingParams
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.multimodal.utils import fetch_image
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
QUESTION = "What is the content of each image?"
|
||||
@@ -1443,6 +1444,8 @@ def run_generate(
|
||||
engine_args.seed = seed
|
||||
if tensor_parallel_size is not None:
|
||||
engine_args.tensor_parallel_size = tensor_parallel_size
|
||||
if current_platform.is_rocm():
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
llm = LLM.from_engine_args(engine_args)
|
||||
|
||||
sampling_params = SamplingParams(
|
||||
@@ -1484,6 +1487,8 @@ def run_chat(
|
||||
engine_args.seed = seed
|
||||
if tensor_parallel_size is not None:
|
||||
engine_args.tensor_parallel_size = tensor_parallel_size
|
||||
if current_platform.is_rocm():
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
llm = LLM.from_engine_args(engine_args)
|
||||
|
||||
sampling_params = (
|
||||
|
||||
@@ -21,6 +21,7 @@ from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset
|
||||
from vllm.lora.request import LoRARequest
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
@@ -2646,6 +2647,8 @@ def main(args):
|
||||
if args.tensor_parallel_size is not None:
|
||||
engine_args.tensor_parallel_size = args.tensor_parallel_size
|
||||
engine_args = maybe_add_vit_cuda_graph_compilation_config(args, engine_args)
|
||||
if current_platform.is_rocm():
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
llm = LLM.from_engine_args(engine_args)
|
||||
|
||||
# Don't want to check the flag multiple times, so just hijack `prompts`.
|
||||
|
||||
@@ -160,3 +160,4 @@ extra_javascript:
|
||||
- https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
|
||||
- mkdocs/javascript/edit_and_feedback.js
|
||||
- mkdocs/javascript/slack_and_forum.js
|
||||
- mkdocs/javascript/reo.js
|
||||
|
||||
+3
-2
@@ -7,7 +7,7 @@ requires = [
|
||||
"setuptools>=77.0.3,<81.0.0",
|
||||
"setuptools-scm>=8.0",
|
||||
"setuptools-rust>=1.9.0",
|
||||
"torch == 2.11.0",
|
||||
"torch == 2.13.0",
|
||||
"wheel",
|
||||
"jinja2",
|
||||
]
|
||||
@@ -122,7 +122,8 @@ python = "./.venv"
|
||||
[tool.typos.files]
|
||||
# these files may be written in non english words
|
||||
extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizers_/*",
|
||||
"benchmarks/sonnet.txt", "tests/lora/data/*", "build/*",
|
||||
"benchmarks/sonnet.txt", "rust/src/bench/src/datasets/sonnet.txt",
|
||||
"tests/lora/data/*", "build/*",
|
||||
"examples/pooling/token_embed/*", "tests/models/language/pooling/*",
|
||||
"vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*",
|
||||
"tests/entrypoints/speech_to_text/transcription/test_transcription_validation.py",
|
||||
|
||||
@@ -4,8 +4,8 @@ packaging>=24.2
|
||||
setuptools==77.0.3 # this version can reuse CMake build dir
|
||||
setuptools-scm>=8
|
||||
setuptools-rust>=1.9.0
|
||||
torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
|
||||
torch==2.11.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
|
||||
torch==2.13.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
|
||||
torch==2.13.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
|
||||
wheel
|
||||
jinja2>=3.1.6
|
||||
regex
|
||||
|
||||
@@ -5,7 +5,7 @@ packaging>=24.2
|
||||
setuptools>=77.0.3,<81.0.0
|
||||
setuptools-scm>=8
|
||||
setuptools-rust>=1.9.0
|
||||
torch==2.11.0
|
||||
torch==2.13.0
|
||||
wheel
|
||||
jinja2>=3.1.6
|
||||
regex
|
||||
|
||||
@@ -6,8 +6,8 @@ setuptools==77.0.3 # this version can reuse CMake build dir
|
||||
numba == 0.65.0; platform_machine != "s390x" # Required for N-gram speculative decoding
|
||||
|
||||
# Dependencies for CPUs
|
||||
torch==2.11.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
|
||||
torch==2.11.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
|
||||
torch==2.13.0+cpu; platform_machine == "x86_64" or platform_machine == "s390x" or platform_machine == "aarch64"
|
||||
torch==2.13.0; platform_system == "Darwin" or platform_machine == "ppc64le" or platform_machine == "riscv64"
|
||||
|
||||
# required for the image processor of minicpm-o-2_6, this must be updated alongside torch
|
||||
torchaudio; platform_machine != "s390x" and platform_machine != "riscv64"
|
||||
|
||||
@@ -4,18 +4,18 @@
|
||||
numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
|
||||
# Dependencies for NVIDIA GPUs
|
||||
torch==2.11.0
|
||||
torch==2.13.0
|
||||
torchaudio==2.11.0
|
||||
# These must be updated alongside torch
|
||||
torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
|
||||
torchvision==0.28.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
|
||||
torchcodec >= 0.14
|
||||
PyNvVideoCodec==2.0.4
|
||||
# FlashInfer should be updated together with the Dockerfile
|
||||
# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
|
||||
# install_requires so the published wheel does not carry an unresolvable pin
|
||||
--extra-index-url https://flashinfer.ai/whl/
|
||||
flashinfer-python==0.6.14
|
||||
flashinfer-cubin==0.6.14
|
||||
flashinfer-python==0.6.15.post1
|
||||
flashinfer-cubin==0.6.15.post1
|
||||
apache-tvm-ffi==0.1.10
|
||||
tilelang==0.1.9
|
||||
nvidia-cudnn-frontend>=1.19.1
|
||||
@@ -26,7 +26,7 @@ fastsafetensors >= 0.3.2
|
||||
|
||||
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
|
||||
nvidia-cutlass-dsl[cu13]==4.6.0
|
||||
quack-kernels>=0.4.0 # Required for tml-fa4
|
||||
quack-kernels>=0.6.1 # Required for CUTLASS DSL 4.6 by MSA
|
||||
|
||||
# Tokenspeed_MLA for faster mla with spec decode
|
||||
tokenspeed-mla==0.1.8; platform_system == "Linux"
|
||||
|
||||
@@ -379,7 +379,7 @@ inflect==5.6.2
|
||||
# via datamodel-code-generator
|
||||
iniconfig==2.0.0
|
||||
# via pytest
|
||||
instanttensor==0.1.5
|
||||
instanttensor==0.1.9
|
||||
# via -r requirements/test/cuda.in
|
||||
interegular==0.3.3
|
||||
# via lm-format-enforcer
|
||||
@@ -1107,7 +1107,7 @@ tokenizers==0.22.2
|
||||
# -r requirements/test/../common.txt
|
||||
# -r requirements/test/cuda.in
|
||||
# transformers
|
||||
torch==2.11.0+cpu
|
||||
torch==2.13.0+cpu
|
||||
# via
|
||||
# -r requirements/test/cuda.in
|
||||
# accelerate
|
||||
@@ -1134,7 +1134,7 @@ torchaudio==2.11.0+cpu
|
||||
# vocos
|
||||
torchcodec==0.14.0+cpu
|
||||
# via -r requirements/test/cuda.in
|
||||
torchvision==0.26.0+cpu
|
||||
torchvision==0.28.0+cpu
|
||||
# via
|
||||
# -r requirements/test/cuda.in
|
||||
# open-clip-torch
|
||||
|
||||
@@ -28,9 +28,9 @@ soundfile # required for audio tests
|
||||
jiwer # required for audio tests
|
||||
tblib # for pickling test exceptions
|
||||
timm >=1.0.17 # required for internvl and gemma3n-mm test
|
||||
torch==2.11.0
|
||||
torch==2.13.0
|
||||
torchaudio==2.11.0
|
||||
torchvision==0.26.0
|
||||
torchvision==0.28.0
|
||||
transformers_stream_generator # required for qwen-vl test
|
||||
matplotlib # required for qwen-vl test
|
||||
mistral_common[image,audio] >= 1.11.5 # required for voxtral test
|
||||
@@ -58,7 +58,7 @@ arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix de
|
||||
numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.3.2
|
||||
instanttensor>=0.1.5; platform_machine == "x86_64"
|
||||
instanttensor>=0.1.9; platform_machine == "x86_64"
|
||||
decord==0.6.0; platform_machine == "x86_64"
|
||||
# terratorch is temporarily disabled while PyPI has the `lightning` package
|
||||
# in `quarantined` status (every published terratorch version transitively
|
||||
|
||||
+12
-10
@@ -159,7 +159,7 @@ cuda-bindings==13.0.3
|
||||
# via torch
|
||||
cuda-pathfinder==1.3.3
|
||||
# via cuda-bindings
|
||||
cuda-toolkit==13.0.2
|
||||
cuda-toolkit==13.0.3.0
|
||||
# via torch
|
||||
cupy-cuda12x==13.6.0
|
||||
# via ray
|
||||
@@ -398,7 +398,7 @@ inflect==5.6.2
|
||||
# via datamodel-code-generator
|
||||
iniconfig==2.0.0
|
||||
# via pytest
|
||||
instanttensor==0.1.5
|
||||
instanttensor==0.1.9
|
||||
# via -r requirements/test/cuda.in
|
||||
interegular==0.3.3
|
||||
# via lm-format-enforcer
|
||||
@@ -599,7 +599,7 @@ numpy==2.2.6
|
||||
# tritonclient
|
||||
# vocos
|
||||
# xgrammar
|
||||
nvidia-cublas==13.1.0.3
|
||||
nvidia-cublas==13.1.1.3
|
||||
# via
|
||||
# cuda-toolkit
|
||||
# nvidia-cudnn-cu13
|
||||
@@ -607,10 +607,12 @@ nvidia-cublas==13.1.0.3
|
||||
nvidia-cuda-cupti==13.0.85
|
||||
# via cuda-toolkit
|
||||
nvidia-cuda-nvrtc==13.0.88
|
||||
# via cuda-toolkit
|
||||
# via
|
||||
# cuda-toolkit
|
||||
# nvidia-cublas
|
||||
nvidia-cuda-runtime==13.0.96
|
||||
# via cuda-toolkit
|
||||
nvidia-cudnn-cu13==9.19.0.56
|
||||
nvidia-cudnn-cu13==9.20.0.48
|
||||
# via torch
|
||||
nvidia-cufft==12.0.0.61
|
||||
# via cuda-toolkit
|
||||
@@ -624,9 +626,9 @@ nvidia-cusparse==12.6.3.3
|
||||
# via
|
||||
# cuda-toolkit
|
||||
# nvidia-cusolver
|
||||
nvidia-cusparselt-cu13==0.8.0
|
||||
nvidia-cusparselt-cu13==0.8.1
|
||||
# via torch
|
||||
nvidia-nccl-cu13==2.28.9
|
||||
nvidia-nccl-cu13==2.29.7
|
||||
# via torch
|
||||
nvidia-nvjitlink==13.0.88
|
||||
# via
|
||||
@@ -1202,7 +1204,7 @@ tokenizers==0.22.2
|
||||
# -r requirements/test/../common.txt
|
||||
# -r requirements/test/cuda.in
|
||||
# transformers
|
||||
torch==2.11.0+cu130
|
||||
torch==2.13.0+cu130
|
||||
# via
|
||||
# -c requirements/cuda.txt
|
||||
# -r requirements/test/cuda.in
|
||||
@@ -1233,7 +1235,7 @@ torchcodec==0.14.0+cu130
|
||||
# via
|
||||
# -c requirements/cuda.txt
|
||||
# -r requirements/test/cuda.in
|
||||
torchvision==0.26.0+cu130
|
||||
torchvision==0.28.0+cu130
|
||||
# via
|
||||
# -c requirements/cuda.txt
|
||||
# -r requirements/test/cuda.in
|
||||
@@ -1270,7 +1272,7 @@ transformers==5.13.1
|
||||
# xgrammar
|
||||
transformers-stream-generator==0.0.5
|
||||
# via -r requirements/test/cuda.in
|
||||
triton==3.6.0
|
||||
triton==3.7.1
|
||||
# via
|
||||
# torch
|
||||
# xgrammar
|
||||
|
||||
@@ -44,5 +44,5 @@ numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
numpy
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.3.2
|
||||
instanttensor>=0.1.5
|
||||
instanttensor>=0.1.9
|
||||
pydantic>=2.12 # 2.11 leads to error on python 3.13
|
||||
|
||||
@@ -54,7 +54,7 @@ arctic-inference==0.1.1 # Required for suffix decoding test
|
||||
numba==0.65.0 # Required for N-gram speculative decoding
|
||||
runai-model-streamer[s3,gcs,azure]==0.15.7
|
||||
fastsafetensors>=0.3.2
|
||||
instanttensor>=0.1.5
|
||||
instanttensor>=0.1.9
|
||||
decord==0.6.0
|
||||
|
||||
# Prithvi tests
|
||||
|
||||
@@ -391,7 +391,7 @@ inflect==7.5.0
|
||||
# via datamodel-code-generator
|
||||
iniconfig==2.3.0
|
||||
# via pytest
|
||||
instanttensor==0.1.6
|
||||
instanttensor==0.1.9
|
||||
# via -r requirements/test/rocm.in
|
||||
interegular==0.3.3
|
||||
# via lm-format-enforcer
|
||||
|
||||
@@ -12,4 +12,4 @@ ray[data]
|
||||
setuptools==78.1.0
|
||||
setuptools-rust>=1.9.0
|
||||
nixl==0.3.0
|
||||
tpu-inference==0.24.0
|
||||
tpu-inference==0.25.0
|
||||
|
||||
@@ -18,4 +18,4 @@ torchvision
|
||||
torchcodec >= 0.14 # Required for the torchcodec video decoding backend
|
||||
|
||||
auto_round_lib==0.14.1
|
||||
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11/vllm_xpu_kernels-0.1.11-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.11.1/vllm_xpu_kernels-0.1.11.1-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
|
||||
Generated
+82
-11
@@ -2585,6 +2585,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
"libm",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3305,6 +3306,16 @@ dependencies = [
|
||||
"getrandom 0.3.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand_distr"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463"
|
||||
dependencies = [
|
||||
"num-traits",
|
||||
"rand 0.9.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rawpointer"
|
||||
version = "0.2.1"
|
||||
@@ -3549,6 +3560,15 @@ dependencies = [
|
||||
"rustc-hash 2.1.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rlimit"
|
||||
version = "0.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f35ee2729c56bb610f6dba436bf78135f728b7373bdffae2ec815b2d3eb98cc3"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rmp"
|
||||
version = "0.8.15"
|
||||
@@ -4855,6 +4875,7 @@ dependencies = [
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -4916,9 +4937,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tonic"
|
||||
version = "0.14.5"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec"
|
||||
checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"axum",
|
||||
@@ -4945,9 +4966,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tonic-build"
|
||||
version = "0.14.5"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734"
|
||||
checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322"
|
||||
dependencies = [
|
||||
"prettyplease",
|
||||
"proc-macro2",
|
||||
@@ -4956,10 +4977,23 @@ dependencies = [
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tonic-prost"
|
||||
version = "0.14.5"
|
||||
name = "tonic-health"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309"
|
||||
checksum = "fcfab99db777fba2802f0dfa861d1628d1ae916fb199d29819941f139ae85082"
|
||||
dependencies = [
|
||||
"prost",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tonic",
|
||||
"tonic-prost",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tonic-prost"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"prost",
|
||||
@@ -4968,9 +5002,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "tonic-prost-build"
|
||||
version = "0.14.5"
|
||||
version = "0.14.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a"
|
||||
checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27"
|
||||
dependencies = [
|
||||
"prettyplease",
|
||||
"proc-macro2",
|
||||
@@ -5439,6 +5473,41 @@ version = "0.9.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
|
||||
|
||||
[[package]]
|
||||
name = "vllm-bench"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"clap",
|
||||
"dirs",
|
||||
"futures",
|
||||
"hf-hub",
|
||||
"image",
|
||||
"indicatif",
|
||||
"mimalloc",
|
||||
"rand 0.9.2",
|
||||
"rand_distr",
|
||||
"rayon",
|
||||
"reqwest 0.12.28",
|
||||
"rlimit",
|
||||
"rustc-hash 1.1.0",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"thiserror-ext",
|
||||
"tiktoken-rs 0.9.1",
|
||||
"tokenizers",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"url",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "vllm-chat"
|
||||
version = "0.1.0"
|
||||
@@ -5507,6 +5576,7 @@ dependencies = [
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"uuid",
|
||||
"vllm-bench",
|
||||
"vllm-chat",
|
||||
"vllm-engine-core-client",
|
||||
"vllm-managed-engine",
|
||||
@@ -5676,6 +5746,7 @@ dependencies = [
|
||||
"tokio-stream",
|
||||
"tokio-util",
|
||||
"tonic",
|
||||
"tonic-health",
|
||||
"tonic-prost",
|
||||
"tonic-prost-build",
|
||||
"tower",
|
||||
@@ -6267,9 +6338,9 @@ checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9"
|
||||
|
||||
[[package]]
|
||||
name = "xgrammar-structural-tag"
|
||||
version = "0.1.0+xgrammar.0.2.2.4d145cc"
|
||||
version = "0.2.0+xgrammar.0.2.4.dd729e7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2436dea2393d55a3b188588aa300c5a8afe8f45a77da52c611fb4498a6c876e6"
|
||||
checksum = "d4d24c842efc3c24e9756aa426d530cbdac0980e49af223cb384e276e981ca0a"
|
||||
dependencies = [
|
||||
"auto_impl",
|
||||
"serde",
|
||||
|
||||
+16
-5
@@ -1,5 +1,6 @@
|
||||
[workspace]
|
||||
members = [
|
||||
"src/bench",
|
||||
"src/chat",
|
||||
"src/cmd",
|
||||
"src/engine-core-client",
|
||||
@@ -32,8 +33,10 @@ base64 = "0.22.1"
|
||||
bytemuck = { version = "1.25.0", features = ["extern_crate_alloc"] }
|
||||
byteorder = "1.5.0"
|
||||
bytes = "1.12.0"
|
||||
chrono = "0.4.42"
|
||||
clap = { version = "4.5.38", features = ["derive", "env"] }
|
||||
criterion = "0.5.1"
|
||||
dirs = "6.0.0"
|
||||
easy-ext = "1.0.3"
|
||||
educe = "0.6.0"
|
||||
enum-as-inner = "0.7.0"
|
||||
@@ -50,7 +53,9 @@ hyper-util = { version = "0.1.20", features = [
|
||||
"service",
|
||||
"tokio",
|
||||
] }
|
||||
image = { version = "0.25.9", default-features = false, features = ["jpeg"] }
|
||||
indexmap = "2.13.0"
|
||||
indicatif = "0.18.4"
|
||||
itertools = "0.14.0"
|
||||
libc = "0.2.177"
|
||||
llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "5390032d6dc8a3e6fdc83acd320260367eb4b9b5", default-features = false, features = ["native-tls"] }
|
||||
@@ -71,10 +76,13 @@ prost-types = "0.14.3"
|
||||
pyo3 = "0.28.3"
|
||||
pythonize = "0.28.0"
|
||||
rand = "0.9.2"
|
||||
rand_distr = "0.5.1"
|
||||
rayon = "1.11.0"
|
||||
reasoning-parser = "1.2.2"
|
||||
reqwest = { version = "0.12.8", default-features = false, features = ["native-tls"] }
|
||||
reqwest-0-13 = { package = "reqwest", version = "0.13.4", default-features = false, features = ["native-tls"] }
|
||||
riptoken = { version = "0.3.0", default-features = false }
|
||||
rlimit = "0.11.0"
|
||||
rmp-serde = "1.3.1"
|
||||
rmpv = { version = "1.3.1", features = ["with-serde"] }
|
||||
rustc-hash = "1.1.0"
|
||||
@@ -110,10 +118,11 @@ tokio = { version = "1.47.1", features = [
|
||||
tokio-openssl = "0.6"
|
||||
tokio-stream = "0.1"
|
||||
tokio-util = { version = "0.7.18", features = ["rt"] }
|
||||
tonic = "0.14.5"
|
||||
tonic-build = "0.14.5"
|
||||
tonic-prost = "0.14.5"
|
||||
tonic-prost-build = "0.14.5"
|
||||
tonic = "0.14.6"
|
||||
tonic-build = "0.14.6"
|
||||
tonic-health = "0.14.6"
|
||||
tonic-prost = "0.14.6"
|
||||
tonic-prost-build = "0.14.6"
|
||||
tool-parser = "1.2.0"
|
||||
tower = { version = "0.5.3", features = ["util"] }
|
||||
tower-http = { version = "0.6.8", features = ["cors", "trace"] }
|
||||
@@ -121,8 +130,10 @@ tracing = { version = "0.1.44", features = ["release_max_level_debug"] }
|
||||
tracing-futures = { version = "0.2.5", features = ["futures-03"] }
|
||||
tracing-subscriber = { version = "0.3.20", features = ["env-filter", "fmt"] }
|
||||
trait-set = "0.3.0"
|
||||
url = "2.5.7"
|
||||
uuid = { version = "1.22.0", features = ["v4"] }
|
||||
validator = { version = "0.20.0", features = ["derive"] }
|
||||
vllm-bench = { path = "src/bench" }
|
||||
vllm-chat = { path = "src/chat" }
|
||||
vllm-engine-core-client = { path = "src/engine-core-client" }
|
||||
vllm-llm = { path = "src/llm" }
|
||||
@@ -133,7 +144,7 @@ vllm-server = { path = "src/server" }
|
||||
vllm-text = { path = "src/text" }
|
||||
vllm-tokenizer = { path = "src/tokenizer" }
|
||||
winnow = { version = "1.0.2", features = ["simd"] }
|
||||
xgrammar-structural-tag = "0.1.0"
|
||||
xgrammar-structural-tag = "0.2.0"
|
||||
zeromq = { version = "0.6.0", default-features = false, features = [
|
||||
"tokio-runtime",
|
||||
"all-transport",
|
||||
|
||||
@@ -14,6 +14,10 @@ service Generate {
|
||||
rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {}
|
||||
}
|
||||
|
||||
service Control {
|
||||
rpc Abort (AbortRequest) returns (AbortResponse) {}
|
||||
}
|
||||
|
||||
// ======================================================================================
|
||||
// Generate Request
|
||||
// ======================================================================================
|
||||
@@ -201,3 +205,12 @@ message TokenIds {
|
||||
repeated uint32 ids = 1;
|
||||
}
|
||||
|
||||
// ======================================================================================
|
||||
// Control
|
||||
// ======================================================================================
|
||||
|
||||
message AbortRequest {
|
||||
repeated string request_ids = 1;
|
||||
}
|
||||
|
||||
message AbortResponse {}
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
# AGENTS.md
|
||||
|
||||
## Project Overview
|
||||
|
||||
Rust rewrite of `vllm bench serve` — a high-performance benchmark client for vLLM serving endpoints. Standalone binary, no Python dependency at runtime.
|
||||
|
||||
Member crate `vllm-bench` of the `rust/` workspace. Uses workspace dependencies and lints; the workspace `[profile.release]` (thin LTO, `panic = "abort"`) applies. Note the workspace bans rustls/ring (`rust/deny.toml`) — all HTTP must stay on native-tls, which is why HF Hub downloads go through `src/hub.rs` (async hf-hub API bridged to sync) instead of hf-hub's ureq backend.
|
||||
|
||||
## Build & Test
|
||||
|
||||
Run from the `rust/` workspace root:
|
||||
|
||||
```bash
|
||||
# Build release binary (rust/target/release/vllm-bench)
|
||||
cargo build -p vllm-bench --release
|
||||
|
||||
# Run all tests
|
||||
cargo test -p vllm-bench
|
||||
|
||||
# Run ignored integration tests (requires network for tokenizer download)
|
||||
cargo test -p vllm-bench -- --ignored
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
- `src/main.rs` — Entry point, mimalloc, tokio runtime, mode dispatch (compare/sweep/multi-run/multi-turn/single)
|
||||
- `src/cli.rs` — clap derive CLI args (~50+ flags)
|
||||
- `src/config.rs` — Validated config from CLI; `GoodputConfig`, `RampUpConfig`, sampling param merging
|
||||
- `src/error.rs` — `BenchError` enum (Http, Json, Tokenizer, Config, EndpointTimeout, Backend, Io)
|
||||
- `src/benchmark.rs` — Core benchmark orchestrator (spawn-per-request with tokio + Semaphore; fetches speculative decoding metrics from `/metrics`)
|
||||
- `src/multi_turn.rs` — Multi-turn conversation orchestrator (channel-based worker pool, sequential turns per conversation)
|
||||
- `src/sweep.rs` — Concurrency/rate parameter sweep (`--sweep-max-concurrency`, `--sweep-request-rate`)
|
||||
- `src/multi_run.rs` — N-run aggregation with mean/std/min/max/CV (`--num-runs`)
|
||||
- `src/compare.rs` — Side-by-side diff of two result JSON files (`--compare`)
|
||||
- `src/tokenizer.rs` — `TokenizerKind` enum: Local(HuggingFace), Tiktoken, OR Server-side `/tokenize`+`/detokenize` fallback
|
||||
- `src/tiktoken.rs` — Tiktoken BPE loader (`.tiktoken`/`.model` files; built-in encodings o200k_base/cl100k_base; pat_str extraction from Python source)
|
||||
- `src/hub.rs` — `HubRepo`: sync facade over hf-hub's async (reqwest/native-tls) API — per-download thread with its own runtime; the sync ureq backend is unusable here because it pulls rustls, which `rust/deny.toml` bans
|
||||
- `src/rate_control.rs` — Gamma/Poisson request scheduling + linear/exponential ramp-up
|
||||
- `src/ready_checker.rs` — Endpoint readiness with retry
|
||||
- `src/backends/` — Backend implementations (enum dispatch, not trait objects)
|
||||
- `mod.rs` — `Backend` enum, `RequestFuncInput`/`RequestFuncOutput` (includes `messages` field for multi-turn)
|
||||
- `streaming.rs` — SSE parser (`StreamedResponseHandler`) with speculative JSON parse for split TCP segments
|
||||
- `openai_completions.rs` — `/v1/completions` backend
|
||||
- `openai_chat.rs` — `/v1/chat/completions` backend (uses `input.messages` when set; zero-copy raw JSON payload for multimodal)
|
||||
- `pooling.rs` — Non-streaming pooling/embedding backends: `openai-embeddings`, `openai-embeddings-chat`, `vllm-pooling`, `vllm-rerank`
|
||||
- `src/datasets/random.rs` — Random dataset generation with rayon parallelism
|
||||
- `src/datasets/random_mm.rs` — Random multimodal dataset (synthetic JPEG images, bucket config sampling, pre-serialized JSON fragments); `--enable-multimodal-chat` pre-builds the chat `messages` array at dataset time (mirrors Python's `apply_multimodal_chat_transformation`)
|
||||
- `src/datasets/sharegpt.rs` — ShareGPT JSON loader + HuggingFace Hub auto-download with caching
|
||||
- `src/datasets/sonnet.rs` and `src/datasets/sonnet.txt` — Sonnet dataset (built-in Shakespeare sonnets via `include_str!("sonnet.txt")`; controllable token length + shared prefix; mirrors Python `SonnetDataset`)
|
||||
- `src/datasets/speed_bench.rs` — NVIDIA SPEED-Bench loader (HF datasets-server API, 6 configs, 11 categories, local cache)
|
||||
- `src/datasets/hf_dataset.rs` — Generic HuggingFace dataset loader (datasets-server API, column auto-detection)
|
||||
- `src/datasets/custom.rs` — Custom JSONL dataset (`{"prompt": ..., "output_tokens": ...}` per line; `--custom-output-len -1` uses per-line output_tokens; prompts always sent raw — no client-side chat template)
|
||||
- `src/datasets/prefix_repetition.rs` — Prefix repetition dataset (N shared prefixes × fresh random suffixes, standard prefix-cache stress; mirrors Python `PrefixRepetitionRandomDataset`)
|
||||
- `src/datasets/random_rerank.rs` — Random rerank dataset (one query + batched documents per request for `vllm-rerank`; `--no-reranker` for embedding-based scoring; mirrors Python `RandomDatasetForReranking`)
|
||||
- `src/datasets/multi_turn.rs` — Multi-turn synthetic generator + ShareGPT multi-turn loader (3-tier prefix sharing: global/conversation/unique-suffix; `per_turn_input_len`)
|
||||
- `src/metrics/mod.rs` — `BenchmarkMetrics` and `MultiTurnMetrics` structs
|
||||
- `src/metrics/calculator.rs` — TTFT/TPOT/ITL/E2EL/throughput stats, goodput SLO checking, peak concurrency, `calculate_multi_turn_metrics`
|
||||
- `src/metrics/steady_state.rs` — Steady-state window detection (in-flight concurrency plateau via two-pointer start/end merge) + plateau throughput/TTFT/TPOT; gated on `--max-concurrency` set + `--request-rate inf` (closed-loop)
|
||||
- `src/output/console.rs` — Terminal output matching Python format + multi-turn per-turn breakdown
|
||||
- `src/output/json.rs` — JSON result file (compatible with Python schema) + multi-turn JSON with `per_turn_metrics`
|
||||
|
||||
## Key Design Decisions
|
||||
|
||||
- **Enum dispatch** for backends (avoids async trait object issues with `dyn`)
|
||||
- **reqwest http1_only()** to match Python aiohttp behavior
|
||||
- **rayon** for parallel dataset generation (key perf win over Python)
|
||||
- **mimalloc** global allocator to reduce contention at 1400+ concurrency (page-agnostic; works on aarch64 64K-page kernels where jemalloc aborts with `LG_PAGE=12` builds)
|
||||
- **Arc\<str\> prompts** zero-copy sharing across tokio tasks (~3GB savings at 100k prompts with 8k-token inputs)
|
||||
- **Spawn-per-request** `tokio::spawn` + `Semaphore` (matches Python asyncio pattern)
|
||||
- **Speculative JSON parse** in SSE handler — detects complete JSON before `\n\n` arrives, improving TTFT/ITL accuracy when TCP segments split
|
||||
- **Tokenizer fallback chain**: Local HF → Tiktoken (`.tiktoken`/`.model` + built-in encodings) → Server-side `/tokenize`+`/detokenize`. Blocking HTTP in rayon threads for server fallback.
|
||||
- **hf-hub** for downloading tokenizers and datasets from HuggingFace Hub
|
||||
- **Pre-serialized mm fragments** (`Arc<str>`) for multimodal: image content stored as JSON strings, zero-copy concatenated into payload — avoids deep-cloning ~200KB+ base64 per request
|
||||
- **Steady-state metrics** (default-on in closed-loop): measure throughput/TTFT/TPOT only over the saturated plateau to cut run-to-run variance at high concurrency; `steady_state` is an `Option` in JSON (`#[serde(default)]` for backward compat), null when the scope gate fails or `--no-steady-state`
|
||||
- **`--prompt-token-ids`** (random dataset only): send token-ID arrays instead of text to skip server-side tokenization; also skips the token-length verification pass (counts exact by construction)
|
||||
- **`--random-range-ratio`** follows Python semantics: lengths sampled uniformly from `[len*(1-r), len*(1+r)]`, default `0.0` = fixed; accepts a float in `[0,1)` or `'{"input": r1, "output": r2}'`. (The pre-2026-07 Rust-only form `[len*r, len]` with default 1.0 is rejected with a migration hint.)
|
||||
- **`prompt_list`** (`Arc<[Arc<str>]>` on `SampleRequest`/`RequestFuncInput`): multiple inputs per request for pooling backends — embeddings batches (`--random-batch-size`) send `"input": [...]`, rerank sends `[0]` as query + `[1..]` as documents
|
||||
- JSON output schema must match Python `vllm bench serve` exactly
|
||||
|
||||
## Common Issues
|
||||
|
||||
- **localhost vs 127.0.0.1**: Some systems resolve `localhost` to IPv6 `::1` while vLLM listens on IPv4 only. Use `127.0.0.1` or the actual hostname.
|
||||
- **Models without tokenizer.json** (e.g., `nvidia/Kimi-K2.5-NVFP4`): Automatically falls back to server-side tokenization. Can also use `--tokenizer` to point to a model with `tokenizer.json`.
|
||||
- **usage.completion_tokens parsing**: vLLM sends final usage chunk with `"choices":[]` (empty array). The usage `if` must be separate from the choices `if` (not `else if`).
|
||||
|
||||
## Typical Usage
|
||||
|
||||
```bash
|
||||
# Embedding benchmark (openai-embeddings, 8 inputs batched per request)
|
||||
./target/release/vllm-bench \
|
||||
--backend openai-embeddings \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model BAAI/bge-large-en-v1.5 \
|
||||
--dataset-name random \
|
||||
--random-input-len 512 \
|
||||
--random-batch-size 8 \
|
||||
--num-prompts 1000 \
|
||||
--save-result
|
||||
|
||||
# vLLM rerank benchmark (one query + 8 documents per request)
|
||||
./target/release/vllm-bench \
|
||||
--backend vllm-rerank \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model BAAI/bge-reranker-v2-m3 \
|
||||
--dataset-name random-rerank \
|
||||
--random-input-len 512 \
|
||||
--random-batch-size 8 \
|
||||
--num-prompts 500 \
|
||||
--save-result
|
||||
|
||||
# Prefix-cache stress (10 shared prefixes, 256+256 tokens)
|
||||
./target/release/vllm-bench \
|
||||
--backend vllm \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model nvidia/Kimi-K2.5-NVFP4 \
|
||||
--dataset-name prefix_repetition \
|
||||
--prefix-repetition-prefix-len 256 \
|
||||
--prefix-repetition-suffix-len 256 \
|
||||
--prefix-repetition-num-prefixes 10 \
|
||||
--num-prompts 1000
|
||||
|
||||
# Custom JSONL workload ({"prompt": ..., "output_tokens": ...} per line)
|
||||
./target/release/vllm-bench \
|
||||
--backend openai-chat \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model nvidia/Kimi-K2.5-NVFP4 \
|
||||
--dataset-name custom \
|
||||
--dataset-path workload.jsonl \
|
||||
--custom-output-len -1 \
|
||||
--num-prompts 1000
|
||||
|
||||
# Random dataset
|
||||
./target/release/vllm-bench \
|
||||
--backend vllm \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model nvidia/Kimi-K2.5-NVFP4 \
|
||||
--dataset-name random \
|
||||
--random-input-len 8192 \
|
||||
--random-output-len 1024 \
|
||||
--ignore-eos \
|
||||
--num-prompts 4096 \
|
||||
--percentile-metrics "ttft,tpot,itl,e2el" \
|
||||
--save-result \
|
||||
--max-concurrency 1400
|
||||
|
||||
# Random multimodal dataset (VLM benchmark)
|
||||
./target/release/vllm-bench \
|
||||
--backend openai-chat \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--dataset-name random-mm \
|
||||
--random-input-len 512 \
|
||||
--random-output-len 128 \
|
||||
--num-prompts 100 \
|
||||
--random-mm-base-items-per-request 1 \
|
||||
--random-mm-limit-mm-per-prompt '{"image": 1, "video": 0}' \
|
||||
--random-mm-bucket-config '{(1024, 800, 1): 1.0}'
|
||||
|
||||
# HuggingFace dataset (WildChat)
|
||||
./target/release/vllm-bench \
|
||||
--backend openai-chat \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model nvidia/Kimi-K2.5-NVFP4 \
|
||||
--dataset-name hf \
|
||||
--dataset-path allenai/WildChat-4.8M \
|
||||
--hf-split train \
|
||||
--num-prompts 1000 \
|
||||
--save-result
|
||||
|
||||
# HuggingFace dataset (LongBench with subset)
|
||||
./target/release/vllm-bench \
|
||||
--backend openai-chat \
|
||||
--base-url http://gb200-10:30000 \
|
||||
--model nvidia/Kimi-K2.5-NVFP4 \
|
||||
--dataset-name hf \
|
||||
--dataset-path THUDM/LongBench \
|
||||
--hf-subset narrativeqa \
|
||||
--hf-split test \
|
||||
--hf-output-len 512 \
|
||||
--num-prompts 200
|
||||
```
|
||||
@@ -0,0 +1 @@
|
||||
@AGENTS.md
|
||||
@@ -0,0 +1,40 @@
|
||||
[package]
|
||||
name = "vllm-bench"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "High-performance benchmark client for vLLM serving endpoints"
|
||||
license.workspace = true
|
||||
|
||||
[dependencies]
|
||||
anyhow.workspace = true
|
||||
base64.workspace = true
|
||||
bytes.workspace = true
|
||||
chrono.workspace = true
|
||||
clap.workspace = true
|
||||
dirs.workspace = true
|
||||
futures.workspace = true
|
||||
hf-hub.workspace = true
|
||||
image.workspace = true
|
||||
indicatif.workspace = true
|
||||
mimalloc.workspace = true
|
||||
rand.workspace = true
|
||||
rand_distr.workspace = true
|
||||
rayon.workspace = true
|
||||
reqwest = { workspace = true, features = ["json", "stream", "http2"] }
|
||||
rlimit.workspace = true
|
||||
rustc-hash.workspace = true
|
||||
serde = { workspace = true, features = ["rc"] }
|
||||
serde_json = { workspace = true, features = ["raw_value"] }
|
||||
thiserror.workspace = true
|
||||
thiserror-ext.workspace = true
|
||||
tiktoken-rs.workspace = true
|
||||
tokenizers.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio-stream.workspace = true
|
||||
tracing.workspace = true
|
||||
tracing-subscriber.workspace = true
|
||||
url.workspace = true
|
||||
uuid.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user