forked from Karylab-cklius/vllm
Compare commits
90
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
07782ff2fc | ||
|
|
0295b4555c | ||
|
|
702e82e422 | ||
|
|
c878b43b64 | ||
|
|
2b84ac669c | ||
|
|
11d3976b88 | ||
|
|
40da9625a1 | ||
|
|
8d9babd4de | ||
|
|
e99ba957ec | ||
|
|
64ac1395e8 | ||
|
|
61cf087680 | ||
|
|
847a57cd12 | ||
|
|
fcd6ac97ed | ||
|
|
95be2a7f22 | ||
|
|
0e60c925cf | ||
|
|
d7ff22204a | ||
|
|
c0bd8b13da | ||
|
|
caeb887bf6 | ||
|
|
6b3166a7c7 | ||
|
|
25e2e136ef | ||
|
|
6874638bc4 | ||
|
|
e24663c5a9 | ||
|
|
c50e105a88 | ||
|
|
a766b30349 | ||
|
|
1faa8cb73c | ||
|
|
e89a91d927 | ||
|
|
909b147197 | ||
|
|
a88b3be7c4 | ||
|
|
a49ea5a58f | ||
|
|
30ebe0dc3c | ||
|
|
cef65f0715 | ||
|
|
6f3b2047ab | ||
|
|
02e8f26cea | ||
|
|
4a00a511bb | ||
|
|
a0d8d944e2 | ||
|
|
df3f537a66 | ||
|
|
7743152957 | ||
|
|
ab33d2a629 | ||
|
|
be3af2d29e | ||
|
|
c656ba3b4d | ||
|
|
dc5fa77a4e | ||
|
|
1e4a084c8e | ||
|
|
7967e854da | ||
|
|
6bd6d0c3c1 | ||
|
|
8e962fef5f | ||
|
|
574fe75245 | ||
|
|
c61a98f529 | ||
|
|
28bffe9466 | ||
|
|
ad65177a19 | ||
|
|
d44a5b6c47 | ||
|
|
1d65283e95 | ||
|
|
c464b57374 | ||
|
|
c5c38e152a | ||
|
|
d00df624f3 | ||
|
|
9752da9d9c | ||
|
|
04925b2202 | ||
|
|
d74278fb67 | ||
|
|
b68fd899d1 | ||
|
|
0b5f9b7204 | ||
|
|
9a8853f781 | ||
|
|
387a1898d9 | ||
|
|
3b30e61507 | ||
|
|
824f9e8f3c | ||
|
|
6cc403e67d | ||
|
|
72d5951d02 | ||
|
|
a3205beffb | ||
|
|
6930becd45 | ||
|
|
03a8770a6d | ||
|
|
bc56a1d56e | ||
|
|
ec7d9e6745 | ||
|
|
3bb4e4311c | ||
|
|
08f8c198ae | ||
|
|
a21cedf4ff | ||
|
|
3ef74cde5d | ||
|
|
cd81cdb399 | ||
|
|
1e828573b4 | ||
|
|
a5ccc85c8c | ||
|
|
b5475d0534 | ||
|
|
9521002f0a | ||
|
|
ec17bdd894 | ||
|
|
bb59c90248 | ||
|
|
5bff999d12 | ||
|
|
bb85929aa6 | ||
|
|
5653021094 | ||
|
|
974d829b05 | ||
|
|
91ac5d9bfd | ||
|
|
23d825aba1 | ||
|
|
f07a128413 | ||
|
|
71cd89264f | ||
|
|
19fab44152 |
@@ -8,7 +8,7 @@ clean_docker_tag() {
|
||||
}
|
||||
|
||||
print_usage_and_exit() {
|
||||
echo "Usage: $0 <registry> <repo> <commit> <branch> <vllm_use_precompiled> <vllm_merge_base_commit> <cache_from> <cache_to>"
|
||||
echo "Usage: $0 <registry> <repo> <commit> <branch> <image_tag> [<image_tag_latest>]"
|
||||
exit 1
|
||||
}
|
||||
|
||||
@@ -142,11 +142,16 @@ resolve_parent_commit() {
|
||||
|
||||
print_bake_config() {
|
||||
echo "--- :page_facing_up: Resolved bake configuration"
|
||||
BAKE_CONFIG_FILE="bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
|
||||
# Write to a temp directory to avoid polluting the repo root (which is the
|
||||
# Docker build context). Files left in the repo root get COPY'd into the
|
||||
# image and can cause duplicate artifact uploads from downstream steps.
|
||||
local bake_tmp
|
||||
bake_tmp="$(mktemp -d)"
|
||||
BAKE_CONFIG_FILE="${bake_tmp}/bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json"
|
||||
docker buildx bake -f "${VLLM_BAKE_FILE_PATH}" -f "${CI_HCL_PATH}" --print "${TARGET}" | tee "${BAKE_CONFIG_FILE}" || true
|
||||
echo "Saved bake config to ${BAKE_CONFIG_FILE}"
|
||||
echo "--- :arrow_down: Uploading bake config to Buildkite"
|
||||
buildkite-agent artifact upload "${BAKE_CONFIG_FILE}"
|
||||
(cd "$(dirname "${BAKE_CONFIG_FILE}")" && buildkite-agent artifact upload "$(basename "${BAKE_CONFIG_FILE}")")
|
||||
}
|
||||
|
||||
#################################
|
||||
@@ -154,7 +159,7 @@ print_bake_config() {
|
||||
#################################
|
||||
print_instance_info
|
||||
|
||||
if [[ $# -lt 7 ]]; then
|
||||
if [[ $# -lt 5 ]]; then
|
||||
print_usage_and_exit
|
||||
fi
|
||||
|
||||
@@ -163,10 +168,8 @@ REGISTRY=$1
|
||||
REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
BRANCH=$4
|
||||
VLLM_USE_PRECOMPILED=$5
|
||||
VLLM_MERGE_BASE_COMMIT=$6
|
||||
IMAGE_TAG=$7
|
||||
IMAGE_TAG_LATEST=${8:-} # only used for main branch, optional
|
||||
IMAGE_TAG=$5
|
||||
IMAGE_TAG_LATEST=${6:-} # only used for main branch, optional
|
||||
|
||||
# build config
|
||||
TARGET="test-ci"
|
||||
@@ -193,8 +196,6 @@ export CACHE_FROM
|
||||
export CACHE_FROM_BASE_BRANCH
|
||||
export CACHE_FROM_MAIN
|
||||
export CACHE_TO
|
||||
export VLLM_USE_PRECOMPILED
|
||||
export VLLM_MERGE_BASE_COMMIT
|
||||
|
||||
# print args
|
||||
echo "--- :mag: Arguments"
|
||||
@@ -202,8 +203,6 @@ echo "REGISTRY: ${REGISTRY}"
|
||||
echo "REPO: ${REPO}"
|
||||
echo "BUILDKITE_COMMIT: ${BUILDKITE_COMMIT}"
|
||||
echo "BRANCH: ${BRANCH}"
|
||||
echo "VLLM_USE_PRECOMPILED: ${VLLM_USE_PRECOMPILED}"
|
||||
echo "VLLM_MERGE_BASE_COMMIT: ${VLLM_MERGE_BASE_COMMIT}"
|
||||
echo "IMAGE_TAG: ${IMAGE_TAG}"
|
||||
echo "IMAGE_TAG_LATEST: ${IMAGE_TAG_LATEST}"
|
||||
|
||||
|
||||
@@ -5,8 +5,7 @@ steps:
|
||||
depends_on: []
|
||||
timeout_in_minutes: 600
|
||||
commands:
|
||||
- if [[ "$BUILDKITE_BRANCH" != "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG; fi
|
||||
- if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $VLLM_USE_PRECOMPILED $VLLM_MERGE_BASE_COMMIT $IMAGE_TAG $IMAGE_TAG_LATEST; fi
|
||||
- if [[ "$BUILDKITE_BRANCH" == "main" ]]; then .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG $IMAGE_TAG_LATEST; else .buildkite/image_build/image_build.sh $REGISTRY $REPO $BUILDKITE_COMMIT $BRANCH $IMAGE_TAG; fi
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: -1 # Agent was lost
|
||||
|
||||
@@ -11,10 +11,10 @@ REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
|
||||
# authenticate with AWS ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
|
||||
|
||||
# skip build if image already exists
|
||||
if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu) ]]; then
|
||||
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then
|
||||
echo "Image not found, proceeding with build..."
|
||||
else
|
||||
echo "Image found"
|
||||
@@ -24,13 +24,13 @@ fi
|
||||
# build
|
||||
docker build --file docker/Dockerfile.cpu \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg buildkite_commit=$BUILDKITE_COMMIT \
|
||||
--build-arg buildkite_commit="$BUILDKITE_COMMIT" \
|
||||
--build-arg VLLM_CPU_AVX512BF16=true \
|
||||
--build-arg VLLM_CPU_AVX512VNNI=true \
|
||||
--build-arg VLLM_CPU_AMXBF16=true \
|
||||
--tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu \
|
||||
--tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu \
|
||||
--target vllm-test \
|
||||
--progress plain .
|
||||
|
||||
# push
|
||||
docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu
|
||||
docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu
|
||||
|
||||
@@ -11,10 +11,10 @@ REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
|
||||
# authenticate with AWS ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
|
||||
|
||||
# skip build if image already exists
|
||||
if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu) ]]; then
|
||||
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu) ]]; then
|
||||
echo "Image not found, proceeding with build..."
|
||||
else
|
||||
echo "Image found"
|
||||
@@ -24,10 +24,10 @@ fi
|
||||
# build
|
||||
docker build --file docker/Dockerfile.cpu \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg buildkite_commit=$BUILDKITE_COMMIT \
|
||||
--tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu \
|
||||
--build-arg buildkite_commit="$BUILDKITE_COMMIT" \
|
||||
--tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu \
|
||||
--target vllm-test \
|
||||
--progress plain .
|
||||
|
||||
# push
|
||||
docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-cpu
|
||||
docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-cpu
|
||||
|
||||
@@ -11,10 +11,10 @@ REPO=$2
|
||||
BUILDKITE_COMMIT=$3
|
||||
|
||||
# authenticate with AWS ECR
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin $REGISTRY
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY"
|
||||
|
||||
# skip build if image already exists
|
||||
if [[ -z $(docker manifest inspect $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu) ]]; then
|
||||
if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu) ]]; then
|
||||
echo "Image not found, proceeding with build..."
|
||||
else
|
||||
echo "Image found"
|
||||
@@ -25,10 +25,10 @@ fi
|
||||
docker build \
|
||||
--file tests/pytorch_ci_hud_benchmark/Dockerfile.hpu \
|
||||
--build-arg max_jobs=16 \
|
||||
--build-arg buildkite_commit=$BUILDKITE_COMMIT \
|
||||
--tag $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu \
|
||||
--build-arg buildkite_commit="$BUILDKITE_COMMIT" \
|
||||
--tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu \
|
||||
--progress plain \
|
||||
https://github.com/vllm-project/vllm-gaudi.git
|
||||
|
||||
# push
|
||||
docker push $REGISTRY/$REPO:$BUILDKITE_COMMIT-hpu
|
||||
docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-hpu
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# We can use this script to compute baseline accuracy on chartqa for vllm.
|
||||
#
|
||||
# Make sure you have lm-eval-harness installed:
|
||||
# pip install "lm-eval[api]>=0.4.9.2"
|
||||
# pip install "lm-eval[api]>=0.4.11"
|
||||
|
||||
usage() {
|
||||
echo``
|
||||
@@ -41,4 +41,4 @@ lm_eval --model vllm-vlm \
|
||||
--tasks chartqa \
|
||||
--batch_size auto \
|
||||
--apply_chat_template \
|
||||
--limit $LIMIT
|
||||
--limit "$LIMIT"
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# We can use this script to compute baseline accuracy on GSM for transformers.
|
||||
#
|
||||
# Make sure you have lm-eval-harness installed:
|
||||
# pip install "lm-eval[api]>=0.4.9.2"
|
||||
# pip install "lm-eval[api]>=0.4.11"
|
||||
|
||||
usage() {
|
||||
echo``
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
# We use this for fp8, which HF does not support.
|
||||
#
|
||||
# Make sure you have lm-eval-harness installed:
|
||||
# pip install "lm-eval[api]>=0.4.9.2"
|
||||
# pip install "lm-eval[api]>=0.4.11"
|
||||
|
||||
usage() {
|
||||
echo``
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
# We use this for fp8, which HF does not support.
|
||||
#
|
||||
# Make sure you have lm-eval-harness installed:
|
||||
# pip install "lm-eval[api]>=0.4.9.2"
|
||||
# pip install "lm-eval[api]>=0.4.11"
|
||||
|
||||
usage() {
|
||||
echo``
|
||||
@@ -20,14 +20,11 @@ usage() {
|
||||
echo
|
||||
}
|
||||
|
||||
while getopts "m:b:l:f:t:" OPT; do
|
||||
while getopts "m:l:f:t:" OPT; do
|
||||
case ${OPT} in
|
||||
m )
|
||||
MODEL="$OPTARG"
|
||||
;;
|
||||
b )
|
||||
BATCH_SIZE="$OPTARG"
|
||||
;;
|
||||
l )
|
||||
LIMIT="$OPTARG"
|
||||
;;
|
||||
|
||||
@@ -15,11 +15,11 @@ DTYPE_FILTER="${DTYPE_FILTER:-}"
|
||||
check_gpus() {
|
||||
if command -v nvidia-smi; then
|
||||
# check the number of GPUs and GPU type.
|
||||
declare -g gpu_count=$(nvidia-smi --list-gpus | wc -l)
|
||||
declare -g gpu_count=$(nvidia-smi --list-gpus | grep -c . || true)
|
||||
elif command -v amd-smi; then
|
||||
declare -g gpu_count=$(amd-smi list | grep 'GPU' | wc -l)
|
||||
declare -g gpu_count=$(amd-smi list | grep -c 'GPU' || true)
|
||||
elif command -v hl-smi; then
|
||||
declare -g gpu_count=$(hl-smi --list | grep -i "Module ID" | wc -l)
|
||||
declare -g gpu_count=$(hl-smi --list | grep -ci "Module ID" || true)
|
||||
fi
|
||||
|
||||
if [[ $gpu_count -gt 0 ]]; then
|
||||
@@ -47,7 +47,7 @@ check_cpus() {
|
||||
declare -g numa_count=$(lscpu | grep "NUMA node(s):" | awk '{print $3}')
|
||||
if [[ $numa_count -gt 0 ]]; then
|
||||
echo "NUMA found."
|
||||
echo $numa_count
|
||||
echo "$numa_count"
|
||||
else
|
||||
echo "Need at least 1 NUMA to run benchmarking."
|
||||
exit 1
|
||||
@@ -434,7 +434,7 @@ run_serving_tests() {
|
||||
|
||||
# iterate over different max_concurrency
|
||||
for max_concurrency in $max_concurrency_list; do
|
||||
new_test_name=$test_name"_qps_"$qps"_concurrency_"$max_concurrency
|
||||
new_test_name="${test_name}_qps_${qps}_concurrency_${max_concurrency}"
|
||||
echo " new test name $new_test_name"
|
||||
# pass the tensor parallel size, the compilation mode, and the optimization
|
||||
# level to the client so that they can be used on the benchmark dashboard
|
||||
@@ -471,7 +471,7 @@ run_serving_tests() {
|
||||
|
||||
# clean up
|
||||
if [[ "${DRY_RUN:-0}" != "1" ]]; then
|
||||
kill -9 $server_pid
|
||||
kill -9 "$server_pid"
|
||||
kill_gpu_processes
|
||||
fi
|
||||
done
|
||||
|
||||
@@ -8,6 +8,21 @@ steps:
|
||||
- group: "Build Python wheels"
|
||||
key: "build-wheels"
|
||||
steps:
|
||||
- label: "Build wheel - aarch64 - CUDA 12.8"
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cuda-12-8
|
||||
agents:
|
||||
queue: arm64_cpu_queue_postmerge
|
||||
commands:
|
||||
# #NOTE: torch_cuda_arch_list is derived from upstream PyTorch build files here:
|
||||
# https://github.com/pytorch/pytorch/blob/main/.ci/aarch64_linux/aarch64_ci_build.sh#L7
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.8.1 --build-arg torch_cuda_arch_list='8.7 8.9 9.0 10.0+PTX' --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"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build wheel - aarch64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
id: build-wheel-arm64-cuda-12-9
|
||||
@@ -51,6 +66,19 @@ steps:
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build wheel - x86_64 - CUDA 12.8"
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cuda-12-8
|
||||
agents:
|
||||
queue: cpu_queue_postmerge
|
||||
commands:
|
||||
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.8.1 --tag vllm-ci:build-image --target build --progress plain -f docker/Dockerfile ."
|
||||
- "mkdir artifacts"
|
||||
- "docker run --rm -v $(pwd)/artifacts:/artifacts_host vllm-ci:build-image bash -c 'cp -r dist /artifacts_host && chmod -R a+rw /artifacts_host'"
|
||||
- "bash .buildkite/scripts/upload-nightly-wheels.sh manylinux_2_31"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
|
||||
- label: "Build wheel - x86_64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
id: build-wheel-x86-cuda-12-9
|
||||
|
||||
@@ -25,7 +25,7 @@ S3_REGION="${AWS_DEFAULT_REGION:-us-west-2}"
|
||||
S3_URL="http://${S3_BUCKET}.s3-website-${S3_REGION}.amazonaws.com"
|
||||
|
||||
# Format ROCm version for path (e.g., "7.1" -> "rocm710")
|
||||
ROCM_VERSION_PATH="rocm$(echo ${ROCM_VERSION} | tr -d '.')"
|
||||
ROCM_VERSION_PATH="rocm$(echo "${ROCM_VERSION}" | tr -d '.')"
|
||||
ROCM_PATH="rocm/${BUILDKITE_COMMIT}/${ROCM_VERSION_PATH}"
|
||||
buildkite-agent annotate --style 'success' --context 'rocm-release-workflow' << EOF
|
||||
## ROCm Wheel and Docker Image Releases
|
||||
|
||||
@@ -83,7 +83,7 @@ case "${1:-}" in
|
||||
exit 1
|
||||
fi
|
||||
|
||||
WHEEL_COUNT=$(ls artifacts/rocm-base-wheels/*.whl 2>/dev/null | wc -l)
|
||||
WHEEL_COUNT=$(find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
|
||||
if [[ "$WHEEL_COUNT" -eq 0 ]]; then
|
||||
echo "ERROR: No wheels found in artifacts/rocm-base-wheels/" >&2
|
||||
exit 1
|
||||
@@ -110,9 +110,9 @@ case "${1:-}" in
|
||||
|
||||
echo ""
|
||||
echo "Downloaded wheels:"
|
||||
ls -lh artifacts/rocm-base-wheels/
|
||||
find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' -exec ls -lh {} \;
|
||||
|
||||
WHEEL_COUNT=$(ls artifacts/rocm-base-wheels/*.whl 2>/dev/null | wc -l)
|
||||
WHEEL_COUNT=$(find artifacts/rocm-base-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
|
||||
echo ""
|
||||
echo "Total: $WHEEL_COUNT wheels"
|
||||
echo "========================================"
|
||||
|
||||
@@ -134,7 +134,7 @@ log_info "Fetching merged PRs from milestone '${MILESTONE}'..."
|
||||
|
||||
# Store PR data in a temp file
|
||||
PR_DATA=$(mktemp)
|
||||
trap "rm -f $PR_DATA" EXIT
|
||||
trap 'rm -f "$PR_DATA"' EXIT
|
||||
|
||||
if ! gh pr list --state merged --search "milestone:${MILESTONE}" \
|
||||
--limit 1000 \
|
||||
|
||||
@@ -27,7 +27,7 @@ function cpu_tests() {
|
||||
podman exec -it "$container_id" bash -c "
|
||||
export TORCH_COMPILE_DISABLE=1
|
||||
set -xve
|
||||
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m" >> $HOME/test_basic.log
|
||||
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m" >> "$HOME"/test_basic.log
|
||||
|
||||
# Run basic model test
|
||||
podman exec -it "$container_id" bash -c "
|
||||
@@ -43,7 +43,7 @@ function cpu_tests() {
|
||||
pytest -v -s tests/models/language/generation/test_common.py::test_models[False-False-5-32-google/gemma-1.1-2b-it]
|
||||
pytest -v -s tests/models/language/pooling/test_classification.py::test_models[float-jason9693/Qwen2.5-1.5B-apeach]
|
||||
# TODO: Below test case tests/models/language/pooling/test_embedding.py::test_models[True-ssmits/Qwen2-7B-Instruct-embed-base] fails on ppc64le. Disabling it for time being.
|
||||
# pytest -v -s tests/models/language/pooling/test_embedding.py -m cpu_model" >> $HOME/test_rest.log
|
||||
# pytest -v -s tests/models/language/pooling/test_embedding.py -m cpu_model" >> "$HOME"/test_rest.log
|
||||
}
|
||||
|
||||
# All of CPU tests are expected to be finished less than 40 mins.
|
||||
|
||||
@@ -16,5 +16,5 @@ echo "--- :docker: Building Docker image"
|
||||
docker build --progress plain --tag "$IMAGE_NAME" --target vllm-test -f docker/Dockerfile.cpu .
|
||||
|
||||
# Run the image, setting --shm-size=4g for tensor parallel.
|
||||
docker run --rm --cpuset-cpus=$CORE_RANGE --cpuset-mems=$NUMA_NODE -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g $IMAGE_NAME \
|
||||
timeout $TIMEOUT_VAL bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
|
||||
docker run --rm --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN -e VLLM_CPU_KVCACHE_SPACE=16 -e VLLM_CPU_CI_ENV=1 -e VLLM_CPU_SIM_MULTI_NUMA=1 --shm-size=4g "$IMAGE_NAME" \
|
||||
timeout "$TIMEOUT_VAL" bash -c "set -euox pipefail; echo \"--- Print packages\"; pip list; echo \"--- Running tests\"; ${TEST_COMMAND}"
|
||||
|
||||
@@ -7,7 +7,7 @@ set -exuo pipefail
|
||||
# Try building the docker image
|
||||
image_name="hpu/upstream-vllm-ci:${BUILDKITE_COMMIT}"
|
||||
container_name="hpu-upstream-vllm-ci-${BUILDKITE_COMMIT}-container"
|
||||
cat <<EOF | docker build -t ${image_name} -f - .
|
||||
cat <<EOF | docker build -t "${image_name}" -f - .
|
||||
FROM gaudi-base-image:latest
|
||||
|
||||
COPY ./ /workspace/vllm
|
||||
@@ -39,12 +39,12 @@ EOF
|
||||
# functions, while other platforms only need one remove_docker_container
|
||||
# function.
|
||||
EXITCODE=1
|
||||
remove_docker_containers() { docker rm -f ${container_name} || true; }
|
||||
remove_docker_containers() { docker rm -f "${container_name}" || true; }
|
||||
trap 'remove_docker_containers; exit $EXITCODE;' EXIT
|
||||
remove_docker_containers
|
||||
|
||||
echo "Running HPU plugin v1 test"
|
||||
docker run --rm --runtime=habana --name=${container_name} --network=host \
|
||||
docker run --rm --runtime=habana --name="${container_name}" --network=host \
|
||||
-e HABANA_VISIBLE_DEVICES=all \
|
||||
-e VLLM_SKIP_WARMUP=true \
|
||||
-e PT_HPU_ENABLE_LAZY_COLLECTIVES=true \
|
||||
|
||||
@@ -41,6 +41,7 @@ get_config() {
|
||||
echo "Error: file '${TEST_RUN_CONFIG_FILE}' does not exist in the warehouse" >&2
|
||||
exit 1
|
||||
fi
|
||||
# shellcheck source=/dev/null
|
||||
source "${TEST_RUN_CONFIG_FILE}"
|
||||
echo "Base docker image name that get from configuration: ${BASE_IMAGE_NAME}"
|
||||
return 0
|
||||
@@ -48,9 +49,8 @@ get_config() {
|
||||
|
||||
# get test running configuration.
|
||||
fetch_vllm_test_cfg
|
||||
get_config
|
||||
# Check if the function call was successful. If not, exit the script.
|
||||
if [ $? -ne 0 ]; then
|
||||
if ! get_config; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
@@ -62,14 +62,14 @@ agent_idx=$(echo "${BUILDKITE_AGENT_NAME}" | awk -F'-' '{print $(NF-1)}')
|
||||
echo "agent_idx: ${agent_idx}"
|
||||
builder_name="cachebuilder${agent_idx}"
|
||||
builder_cache_dir="/mnt/docker-cache${agent_idx}"
|
||||
mkdir -p ${builder_cache_dir}
|
||||
mkdir -p "${builder_cache_dir}"
|
||||
|
||||
# Try building the docker image
|
||||
cat <<EOF | DOCKER_BUILDKIT=1 docker build \
|
||||
--add-host cache-service-vllm.nginx-pypi-cache.svc.cluster.local:${PYPI_CACHE_HOST} \
|
||||
--builder ${builder_name} --cache-from type=local,src=${builder_cache_dir} \
|
||||
--cache-to type=local,dest=${builder_cache_dir},mode=max \
|
||||
--progress=plain --load -t ${image_name} -f - .
|
||||
--add-host cache-service-vllm.nginx-pypi-cache.svc.cluster.local:"${PYPI_CACHE_HOST}" \
|
||||
--builder "${builder_name}" --cache-from type=local,src="${builder_cache_dir}" \
|
||||
--cache-to type=local,dest="${builder_cache_dir}",mode=max \
|
||||
--progress=plain --load -t "${image_name}" -f - .
|
||||
FROM ${BASE_IMAGE_NAME}
|
||||
|
||||
# Define environments
|
||||
@@ -116,7 +116,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/$(uname -i)-linux/devlib && \
|
||||
python3 -m pip install -v -e /workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
@@ -139,7 +139,7 @@ trap remove_docker_container EXIT
|
||||
# Generate corresponding --device args based on BUILDKITE_AGENT_NAME
|
||||
# Ascend NPU BUILDKITE_AGENT_NAME format is {hostname}-{agent_idx}-{npu_card_num}cards, and agent_idx starts from 1.
|
||||
# e.g. atlas-a2-001-1-2cards means this is the 1-th agent on atlas-a2-001 host, and it has 2 NPU cards.
|
||||
# returns --device /dev/davinci0 --device /dev/davinci1
|
||||
# returns one argument per line: --device, /dev/davinciX, ...
|
||||
parse_and_gen_devices() {
|
||||
local input="$1"
|
||||
local index cards_num
|
||||
@@ -151,29 +151,24 @@ parse_and_gen_devices() {
|
||||
return 1
|
||||
fi
|
||||
|
||||
local devices=""
|
||||
local i=0
|
||||
while (( i < cards_num )); do
|
||||
local dev_idx=$(((index - 1)*cards_num + i ))
|
||||
devices="$devices --device /dev/davinci${dev_idx}"
|
||||
printf '%s\n' "--device"
|
||||
printf '%s\n' "/dev/davinci${dev_idx}"
|
||||
((i++))
|
||||
done
|
||||
|
||||
# trim leading space
|
||||
devices="${devices#"${devices%%[![:space:]]*}"}"
|
||||
# Output devices: assigned to the caller variable
|
||||
printf '%s' "$devices"
|
||||
}
|
||||
|
||||
devices=$(parse_and_gen_devices "${BUILDKITE_AGENT_NAME}") || exit 1
|
||||
mapfile -t device_args < <(parse_and_gen_devices "${BUILDKITE_AGENT_NAME}") || exit 1
|
||||
|
||||
# Run the image and execute the Out-Of-Tree (OOT) platform interface test case on Ascend NPU hardware.
|
||||
# This test checks whether the OOT platform interface is functioning properly in conjunction with
|
||||
# the hardware plugin vllm-ascend.
|
||||
model_cache_dir=/mnt/modelscope${agent_idx}
|
||||
mkdir -p ${model_cache_dir}
|
||||
mkdir -p "${model_cache_dir}"
|
||||
docker run \
|
||||
${devices} \
|
||||
"${device_args[@]}" \
|
||||
--device /dev/davinci_manager \
|
||||
--device /dev/devmm_svm \
|
||||
--device /dev/hisi_hdc \
|
||||
@@ -182,7 +177,7 @@ docker run \
|
||||
-v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
|
||||
-v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
|
||||
-v /etc/ascend_install.info:/etc/ascend_install.info \
|
||||
-v ${model_cache_dir}:/root/.cache/modelscope \
|
||||
-v "${model_cache_dir}":/root/.cache/modelscope \
|
||||
--entrypoint="" \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
|
||||
@@ -61,7 +61,7 @@ echo "Results will be stored in: $RESULTS_DIR"
|
||||
echo "--- Installing Python dependencies ---"
|
||||
python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
|
||||
&& python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
|
||||
&& python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.9.2" \
|
||||
&& python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
|
||||
&& python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
|
||||
echo "--- Python dependencies installed ---"
|
||||
|
||||
|
||||
@@ -61,7 +61,7 @@ echo "Results will be stored in: $RESULTS_DIR"
|
||||
echo "--- Installing Python dependencies ---"
|
||||
python3 -m pip install --progress-bar off git+https://github.com/thuml/depyf.git \
|
||||
&& python3 -m pip install --progress-bar off pytest pytest-asyncio tpu-info \
|
||||
&& python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.9.2" \
|
||||
&& python3 -m pip install --progress-bar off "lm-eval[api]>=0.4.11" \
|
||||
&& python3 -m pip install --progress-bar off hf-transfer tblib==3.1.0
|
||||
echo "--- Python dependencies installed ---"
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ image_name="xpu/vllm-ci:${BUILDKITE_COMMIT}"
|
||||
container_name="xpu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
|
||||
|
||||
# Try building the docker image
|
||||
docker build -t ${image_name} -f docker/Dockerfile.xpu .
|
||||
docker build -t "${image_name}" -f docker/Dockerfile.xpu .
|
||||
|
||||
# Setup cleanup
|
||||
remove_docker_container() {
|
||||
|
||||
@@ -21,16 +21,16 @@ echo "Pushing original tag $ORIG_TAG_NAME$ORIG_TAG_SUFFIX to new nightly tag nam
|
||||
|
||||
# pull original arch-dependent images from AWS ECR Public
|
||||
aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-x86_64$ORIG_TAG_SUFFIX
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-aarch64$ORIG_TAG_SUFFIX
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-x86_64"$ORIG_TAG_SUFFIX"
|
||||
docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-aarch64"$ORIG_TAG_SUFFIX"
|
||||
# tag arch-dependent images
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-x86_64$ORIG_TAG_SUFFIX vllm/vllm-openai:$TAG_NAME-x86_64
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$ORIG_TAG_NAME-aarch64$ORIG_TAG_SUFFIX vllm/vllm-openai:$TAG_NAME-aarch64
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-x86_64"$ORIG_TAG_SUFFIX" vllm/vllm-openai:"$TAG_NAME"-x86_64
|
||||
docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:"$ORIG_TAG_NAME"-aarch64"$ORIG_TAG_SUFFIX" vllm/vllm-openai:"$TAG_NAME"-aarch64
|
||||
# push arch-dependent images to DockerHub
|
||||
docker push vllm/vllm-openai:$TAG_NAME-x86_64
|
||||
docker push vllm/vllm-openai:$TAG_NAME-aarch64
|
||||
docker push vllm/vllm-openai:"$TAG_NAME"-x86_64
|
||||
docker push vllm/vllm-openai:"$TAG_NAME"-aarch64
|
||||
# push arch-independent manifest to DockerHub
|
||||
docker manifest create vllm/vllm-openai:$TAG_NAME vllm/vllm-openai:$TAG_NAME-x86_64 vllm/vllm-openai:$TAG_NAME-aarch64 --amend
|
||||
docker manifest create vllm/vllm-openai:$TAG_NAME-$BUILDKITE_COMMIT vllm/vllm-openai:$TAG_NAME-x86_64 vllm/vllm-openai:$TAG_NAME-aarch64 --amend
|
||||
docker manifest push vllm/vllm-openai:$TAG_NAME
|
||||
docker manifest push vllm/vllm-openai:$TAG_NAME-$BUILDKITE_COMMIT
|
||||
docker manifest create vllm/vllm-openai:"$TAG_NAME" vllm/vllm-openai:"$TAG_NAME"-x86_64 vllm/vllm-openai:"$TAG_NAME"-aarch64 --amend
|
||||
docker manifest create vllm/vllm-openai:"$TAG_NAME"-"$BUILDKITE_COMMIT" vllm/vllm-openai:"$TAG_NAME"-x86_64 vllm/vllm-openai:"$TAG_NAME"-aarch64 --amend
|
||||
docker manifest push vllm/vllm-openai:"$TAG_NAME"
|
||||
docker manifest push vllm/vllm-openai:"$TAG_NAME"-"$BUILDKITE_COMMIT"
|
||||
|
||||
@@ -29,7 +29,7 @@ fi
|
||||
if ! command -v uv &> /dev/null; then
|
||||
echo "Installing UV package manager..."
|
||||
curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
source $HOME/.local/bin/env
|
||||
source "$HOME"/.local/bin/env
|
||||
fi
|
||||
|
||||
# Clone Prime-RL repository at specific branch for reproducible tests
|
||||
|
||||
@@ -51,14 +51,14 @@ for BACK in "${BACKENDS[@]}"; do
|
||||
--enable-eplb \
|
||||
--trust-remote-code \
|
||||
--max-model-len 2048 \
|
||||
--all2all-backend $BACK \
|
||||
--port $PORT &
|
||||
--all2all-backend "$BACK" \
|
||||
--port "$PORT" &
|
||||
SERVER_PID=$!
|
||||
wait_for_server $PORT
|
||||
wait_for_server "$PORT"
|
||||
|
||||
TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
|
||||
OUT="${OUT_DIR}/${TAG}_${BACK}.json"
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port $PORT --num-questions ${NUM_Q} --save-results ${OUT}
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
|
||||
python3 - <<PY
|
||||
import json; acc=json.load(open('${OUT}'))['accuracy']
|
||||
print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")
|
||||
|
||||
@@ -47,20 +47,20 @@ for BACK in "${BACKENDS[@]}"; do
|
||||
vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--enable-eplb \
|
||||
--all2all-backend $BACK \
|
||||
--all2all-backend "$BACK" \
|
||||
--eplb-config '{"window_size":10, "step_interval":100, "num_redundant_experts":0, "log_balancedness":true}' \
|
||||
--tensor-parallel-size ${TENSOR_PARALLEL_SIZE} \
|
||||
--data-parallel-size ${DATA_PARALLEL_SIZE} \
|
||||
--tensor-parallel-size "${TENSOR_PARALLEL_SIZE}" \
|
||||
--data-parallel-size "${DATA_PARALLEL_SIZE}" \
|
||||
--enable-expert-parallel \
|
||||
--trust-remote-code \
|
||||
--max-model-len 2048 \
|
||||
--port $PORT &
|
||||
--port "$PORT" &
|
||||
SERVER_PID=$!
|
||||
wait_for_server $PORT
|
||||
wait_for_server "$PORT"
|
||||
|
||||
TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
|
||||
OUT="${OUT_DIR}/${TAG}_${BACK}.json"
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port $PORT --num-questions ${NUM_Q} --save-results ${OUT}
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
|
||||
python3 - <<PY
|
||||
import json; acc=json.load(open('${OUT}'))['accuracy']
|
||||
print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")
|
||||
|
||||
@@ -51,20 +51,20 @@ for BACK in "${BACKENDS[@]}"; do
|
||||
--tensor-parallel-size 4 \
|
||||
--enable-expert-parallel \
|
||||
--enable-eplb \
|
||||
--all2all-backend $BACK \
|
||||
--all2all-backend "$BACK" \
|
||||
--eplb-config '{"window_size":200,"step_interval":600,"use_async":true}' \
|
||||
--speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":1}' \
|
||||
--trust-remote-code \
|
||||
--max-model-len 2048 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
"${PLATFORM_ARGS[@]}" \
|
||||
--port $PORT &
|
||||
--port "$PORT" &
|
||||
SERVER_PID=$!
|
||||
wait_for_server $PORT
|
||||
wait_for_server "$PORT"
|
||||
|
||||
TAG=$(echo "$MODEL" | tr '/: \\n' '_____')
|
||||
OUT="${OUT_DIR}/${TAG}_${BACK}.json"
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port $PORT --num-questions ${NUM_Q} --save-results ${OUT}
|
||||
python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}"
|
||||
python3 - <<PY
|
||||
import json; acc=json.load(open('${OUT}'))['accuracy']
|
||||
print(f"${MODEL} ${BACK}: accuracy {acc:.3f}")
|
||||
|
||||
@@ -9,10 +9,11 @@ ENV_FILE=$1
|
||||
|
||||
# For testing on local vm, use `set -a` to export all variables
|
||||
source /etc/environment
|
||||
source $ENV_FILE
|
||||
# shellcheck source=/dev/null
|
||||
source "$ENV_FILE"
|
||||
|
||||
remove_docker_container() {
|
||||
docker rm -f $CONTAINER_NAME || true;
|
||||
docker rm -f "$CONTAINER_NAME" || true;
|
||||
}
|
||||
|
||||
trap remove_docker_container EXIT
|
||||
@@ -41,13 +42,13 @@ echo
|
||||
echo "starting docker...$CONTAINER_NAME"
|
||||
echo
|
||||
docker run \
|
||||
-v $DOWNLOAD_DIR:$DOWNLOAD_DIR \
|
||||
--env-file $ENV_FILE \
|
||||
-v "$DOWNLOAD_DIR":"$DOWNLOAD_DIR" \
|
||||
--env-file "$ENV_FILE" \
|
||||
-e HF_TOKEN="$HF_TOKEN" \
|
||||
-e TARGET_COMMIT=$BUILDKITE_COMMIT \
|
||||
-e MODEL=$MODEL \
|
||||
-e TARGET_COMMIT="$BUILDKITE_COMMIT" \
|
||||
-e MODEL="$MODEL" \
|
||||
-e WORKSPACE=/workspace \
|
||||
--name $CONTAINER_NAME \
|
||||
--name "$CONTAINER_NAME" \
|
||||
-d \
|
||||
--privileged \
|
||||
--network host \
|
||||
|
||||
@@ -42,21 +42,21 @@ echo "lanching vllm..."
|
||||
echo "logging to $VLLM_LOG"
|
||||
echo
|
||||
|
||||
vllm serve $MODEL \
|
||||
vllm serve "$MODEL" \
|
||||
--seed 42 \
|
||||
--max-num-seqs $MAX_NUM_SEQS \
|
||||
--max-num-batched-tokens $MAX_NUM_BATCHED_TOKENS \
|
||||
--tensor-parallel-size $TENSOR_PARALLEL_SIZE \
|
||||
--max-num-seqs "$MAX_NUM_SEQS" \
|
||||
--max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \
|
||||
--tensor-parallel-size "$TENSOR_PARALLEL_SIZE" \
|
||||
--no-enable-prefix-caching \
|
||||
--download_dir $DOWNLOAD_DIR \
|
||||
--max-model-len $MAX_MODEL_LEN > "$VLLM_LOG" 2>&1 &
|
||||
--download_dir "$DOWNLOAD_DIR" \
|
||||
--max-model-len "$MAX_MODEL_LEN" > "$VLLM_LOG" 2>&1 &
|
||||
|
||||
|
||||
echo "wait for 20 minutes.."
|
||||
echo
|
||||
# sleep 1200
|
||||
# wait for 10 minutes...
|
||||
for i in {1..120}; do
|
||||
for _ in {1..120}; do
|
||||
# TODO: detect other type of errors.
|
||||
if grep -Fq "raise RuntimeError" "$VLLM_LOG"; then
|
||||
echo "Detected RuntimeError, exiting."
|
||||
@@ -78,11 +78,11 @@ echo "logging to $BM_LOG"
|
||||
echo
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--dataset-name sonnet \
|
||||
--dataset-path benchmarks/sonnet_4x.txt \
|
||||
--sonnet-input-len $INPUT_LEN \
|
||||
--sonnet-output-len $OUTPUT_LEN \
|
||||
--sonnet-input-len "$INPUT_LEN" \
|
||||
--sonnet-output-len "$OUTPUT_LEN" \
|
||||
--ignore-eos > "$BM_LOG"
|
||||
|
||||
echo "completed..."
|
||||
|
||||
@@ -76,16 +76,15 @@ mkdir -p "$INDICES_OUTPUT_DIR"
|
||||
# this indices have relative paths that could work as long as it is next to the wheel directory in s3
|
||||
# i.e., the wheels are always in s3://vllm-wheels/<commit>/
|
||||
# and indices can be placed in /<commit>/, or /nightly/, or /<version>/
|
||||
if [[ ! -z "$DEFAULT_VARIANT_ALIAS" ]]; then
|
||||
alias_arg="--alias-to-default $DEFAULT_VARIANT_ALIAS"
|
||||
else
|
||||
alias_arg=""
|
||||
alias_args=()
|
||||
if [[ -n "$DEFAULT_VARIANT_ALIAS" ]]; then
|
||||
alias_args=(--alias-to-default "$DEFAULT_VARIANT_ALIAS")
|
||||
fi
|
||||
|
||||
# HACK: we do not need regex module here, but it is required by pre-commit hook
|
||||
# To avoid any external dependency, we simply replace it back to the stdlib re module
|
||||
sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
|
||||
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" $alias_arg
|
||||
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "commit $BUILDKITE_COMMIT" "${alias_args[@]}"
|
||||
|
||||
# copy indices to /<commit>/ unconditionally
|
||||
echo "Uploading indices to $S3_COMMIT_PREFIX"
|
||||
@@ -100,9 +99,9 @@ fi
|
||||
# re-generate and copy to /<pure_version>/ only if it does not have "dev" in the version
|
||||
if [[ "$version" != *"dev"* ]]; then
|
||||
echo "Re-generating indices for /$pure_version/"
|
||||
rm -rf "$INDICES_OUTPUT_DIR/*"
|
||||
rm -rf "${INDICES_OUTPUT_DIR:?}/*"
|
||||
mkdir -p "$INDICES_OUTPUT_DIR"
|
||||
# wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
|
||||
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --wheel-dir "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" $alias_arg
|
||||
$PYTHON .buildkite/scripts/generate-nightly-index.py --version "$pure_version" --wheel-dir "$SUBPATH" --current-objects "$obj_json" --output-dir "$INDICES_OUTPUT_DIR" --comment "version $pure_version" "${alias_args[@]}"
|
||||
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/$pure_version/"
|
||||
fi
|
||||
|
||||
@@ -7,7 +7,7 @@ SUBPATH=$BUILDKITE_COMMIT
|
||||
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
|
||||
|
||||
RELEASE_VERSION=$(buildkite-agent meta-data get release-version)
|
||||
GIT_VERSION=$(git describe --exact-match --tags $BUILDKITE_COMMIT 2>/dev/null)
|
||||
GIT_VERSION=$(git describe --exact-match --tags "$BUILDKITE_COMMIT" 2>/dev/null)
|
||||
|
||||
echo "Release version from Buildkite: $RELEASE_VERSION"
|
||||
|
||||
@@ -55,7 +55,7 @@ mkdir -p $DIST_DIR
|
||||
aws s3 cp --recursive --exclude "*" --include "vllm-${PURE_VERSION}*.whl" --exclude "*dev*" --exclude "*rc[0-9]*" "$S3_COMMIT_PREFIX" $DIST_DIR
|
||||
echo "Wheels copied to local directory"
|
||||
# generate source tarball
|
||||
git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" $BUILDKITE_COMMIT
|
||||
git archive --format=tar.gz --output="$DIST_DIR/vllm-${PURE_VERSION}.tar.gz" "$BUILDKITE_COMMIT"
|
||||
ls -la $DIST_DIR
|
||||
|
||||
# upload wheels to PyPI (only default variant, i.e. files without '+' in the name)
|
||||
@@ -65,6 +65,6 @@ if [[ -z "$PYPI_WHEEL_FILES" ]]; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
python3 -m twine check $PYPI_WHEEL_FILES
|
||||
python3 -m twine upload --non-interactive --verbose $PYPI_WHEEL_FILES
|
||||
python3 -m twine check "$PYPI_WHEEL_FILES"
|
||||
python3 -m twine upload --non-interactive --verbose "$PYPI_WHEEL_FILES"
|
||||
echo "Wheels uploaded to PyPI"
|
||||
|
||||
@@ -55,7 +55,7 @@ mkdir -p all-rocm-wheels
|
||||
cp artifacts/rocm-base-wheels/*.whl all-rocm-wheels/ 2>/dev/null || true
|
||||
cp artifacts/rocm-vllm-wheel/*.whl all-rocm-wheels/ 2>/dev/null || true
|
||||
|
||||
WHEEL_COUNT=$(ls all-rocm-wheels/*.whl 2>/dev/null | wc -l)
|
||||
WHEEL_COUNT=$(find all-rocm-wheels -maxdepth 1 -name '*.whl' 2>/dev/null | wc -l)
|
||||
echo "Total wheels to upload: $WHEEL_COUNT"
|
||||
|
||||
if [ "$WHEEL_COUNT" -eq 0 ]; then
|
||||
@@ -115,7 +115,7 @@ if [[ "$BUILDKITE_BRANCH" == "main" && "$BUILDKITE_PULL_REQUEST" == "false" ]] |
|
||||
fi
|
||||
|
||||
# Extract version from vLLM wheel and update version-specific index
|
||||
VLLM_WHEEL=$(ls all-rocm-wheels/vllm*.whl 2>/dev/null | head -1)
|
||||
VLLM_WHEEL=$(find all-rocm-wheels -maxdepth 1 -name 'vllm*.whl' 2>/dev/null | head -1)
|
||||
if [ -n "$VLLM_WHEEL" ]; then
|
||||
VERSION=$(unzip -p "$VLLM_WHEEL" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
|
||||
echo "Version in wheel: $VERSION"
|
||||
|
||||
+1669
-1
File diff suppressed because it is too large
Load Diff
@@ -51,9 +51,11 @@ steps:
|
||||
mirror_hardwares: [amdexperimental]
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/detokenizer
|
||||
- tests/multimodal
|
||||
- tests/utils_
|
||||
commands:
|
||||
- pytest -v -s detokenizer
|
||||
- pytest -v -s -m 'not cpu_test' multimodal
|
||||
- pytest -v -s utils_
|
||||
|
||||
|
||||
@@ -165,6 +165,7 @@ steps:
|
||||
num_devices: 2
|
||||
num_nodes: 2
|
||||
no_plugin: true
|
||||
optional: true # TODO: revert once infra issue solved
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
@@ -197,6 +198,17 @@ steps:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- DP_EP=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: CrossLayer KV layout Distributed NixlConnector PD accuracy tests (4 GPUs)
|
||||
timeout_in_minutes: 30
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors.txt
|
||||
- CROSS_LAYERS_BLOCKS=True bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Pipeline + Context Parallelism (4 GPUs))
|
||||
timeout_in_minutes: 60
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
|
||||
@@ -108,9 +108,11 @@ steps:
|
||||
timeout_in_minutes: 50
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/detokenizer
|
||||
- tests/multimodal
|
||||
- tests/utils_
|
||||
commands:
|
||||
- pytest -v -s detokenizer
|
||||
- pytest -v -s -m 'not cpu_test' multimodal
|
||||
- pytest -v -s utils_
|
||||
|
||||
@@ -123,6 +125,7 @@ steps:
|
||||
- tests/test_inputs.py
|
||||
- tests/test_outputs.py
|
||||
- tests/test_pooling_params.py
|
||||
- tests/test_ray_env.py
|
||||
- tests/multimodal
|
||||
- tests/renderers
|
||||
- tests/standalone_tests/lazy_imports.py
|
||||
@@ -136,6 +139,7 @@ steps:
|
||||
- pytest -v -s test_inputs.py
|
||||
- pytest -v -s test_outputs.py
|
||||
- pytest -v -s test_pooling_params.py
|
||||
- pytest -v -s test_ray_env.py
|
||||
- pytest -v -s -m 'cpu_test' multimodal
|
||||
- pytest -v -s renderers
|
||||
- pytest -v -s tokenizers_
|
||||
|
||||
@@ -18,4 +18,4 @@ steps:
|
||||
depends_on:
|
||||
- image-build-amd
|
||||
commands:
|
||||
- pytest -v -s -m 'not skip_v1' samplers
|
||||
- pytest -v -s -m samplers
|
||||
|
||||
@@ -771,6 +771,25 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# DeepSeek V3 fused A GEMM kernel (requires SM 9.0+, Hopper and later)
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0f;11.0f" "${CUDA_ARCHS}")
|
||||
else()
|
||||
cuda_archs_loose_intersection(DSV3_FUSED_A_GEMM_ARCHS "9.0a;10.0a;10.1a;10.3a" "${CUDA_ARCHS}")
|
||||
endif()
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSV3_FUSED_A_GEMM_ARCHS)
|
||||
set(DSV3_FUSED_A_GEMM_SRC "csrc/dsv3_fused_a_gemm.cu")
|
||||
set_gencode_flags_for_srcs(
|
||||
SRCS "${DSV3_FUSED_A_GEMM_SRC}"
|
||||
CUDA_ARCHS "${DSV3_FUSED_A_GEMM_ARCHS}")
|
||||
list(APPEND VLLM_EXT_SRC ${DSV3_FUSED_A_GEMM_SRC})
|
||||
list(APPEND VLLM_GPU_FLAGS "-DENABLE_DSV3_FUSED_A_GEMM=1")
|
||||
message(STATUS "Building dsv3_fused_a_gemm for archs: ${DSV3_FUSED_A_GEMM_ARCHS}")
|
||||
else()
|
||||
message(STATUS "Not building dsv3_fused_a_gemm as no compatible archs found "
|
||||
"in CUDA target architectures.")
|
||||
endif()
|
||||
|
||||
# moe_data.cu is used by all CUTLASS MoE kernels.
|
||||
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
|
||||
cuda_archs_loose_intersection(CUTLASS_MOE_DATA_ARCHS "9.0a;10.0f;11.0f;12.0f" "${CUDA_ARCHS}")
|
||||
|
||||
@@ -46,10 +46,10 @@ echo "VLLM_LOGGING_LEVEL=$VLLM_LOGGING_LEVEL"
|
||||
echo "RESULT_FILE=$RESULT"
|
||||
echo "====================== AUTO TUNEPARAMETERS ===================="
|
||||
|
||||
rm -rf $LOG_FOLDER
|
||||
rm -rf $PROFILE_PATH
|
||||
mkdir -p $LOG_FOLDER
|
||||
mkdir -p $PROFILE_PATH
|
||||
rm -rf "$LOG_FOLDER"
|
||||
rm -rf "$PROFILE_PATH"
|
||||
mkdir -p "$LOG_FOLDER"
|
||||
mkdir -p "$PROFILE_PATH"
|
||||
|
||||
cd "$BASE/vllm"
|
||||
|
||||
@@ -114,7 +114,7 @@ start_server() {
|
||||
|
||||
# wait for 10 minutes...
|
||||
server_started=0
|
||||
for i in {1..60}; do
|
||||
for _ in {1..60}; do
|
||||
# This line checks whether the server is still alive or not,
|
||||
# since that we should always have permission to send signal to the server process.
|
||||
kill -0 $server_pid 2> /dev/null || break
|
||||
@@ -145,12 +145,12 @@ run_benchmark() {
|
||||
local vllm_log="$LOG_FOLDER/vllm_log_${max_num_seqs}_${max_num_batched_tokens}.txt"
|
||||
echo "vllm_log: $vllm_log"
|
||||
echo
|
||||
rm -f $vllm_log
|
||||
rm -f "$vllm_log"
|
||||
pkill -if "vllm serve" || true
|
||||
|
||||
echo "starting server..."
|
||||
# Call start_server without a profile_dir to avoid profiling overhead
|
||||
start_server $gpu_memory_utilization $max_num_seqs $max_num_batched_tokens $vllm_log ""
|
||||
start_server "$gpu_memory_utilization" "$max_num_seqs" "$max_num_batched_tokens" "$vllm_log" ""
|
||||
result=$?
|
||||
if [[ "$result" -eq 1 ]]; then
|
||||
echo "server failed to start. gpu_memory_utilization:$gpu_memory_utilization, max_num_seqs:$max_num_seqs, max_num_batched_tokens: $max_num_batched_tokens"
|
||||
@@ -168,15 +168,15 @@ run_benchmark() {
|
||||
# --profile flag is removed from this call
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--dataset-name random \
|
||||
--random-input-len $adjusted_input_len \
|
||||
--random-output-len $OUTPUT_LEN \
|
||||
--random-output-len "$OUTPUT_LEN" \
|
||||
--ignore-eos \
|
||||
--disable-tqdm \
|
||||
--request-rate inf \
|
||||
--percentile-metrics ttft,tpot,itl,e2el \
|
||||
--goodput e2el:$MAX_LATENCY_ALLOWED_MS \
|
||||
--goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \
|
||||
--num-prompts 1000 \
|
||||
--random-prefix-len $prefix_len \
|
||||
--host "$HOSTNAME" \
|
||||
@@ -195,20 +195,20 @@ run_benchmark() {
|
||||
request_rate=$((${throughput%.*} + 1))
|
||||
while ((request_rate > 0)); do
|
||||
# clear prefix cache
|
||||
curl -X POST http://${HOSTNAME}:8004/reset_prefix_cache
|
||||
curl -X POST http://"${HOSTNAME}":8004/reset_prefix_cache
|
||||
sleep 5
|
||||
bm_log="$LOG_FOLDER/bm_log_${max_num_seqs}_${max_num_batched_tokens}_requestrate_${request_rate}.txt"
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--dataset-name random \
|
||||
--random-input-len $adjusted_input_len \
|
||||
--random-output-len $OUTPUT_LEN \
|
||||
--random-output-len "$OUTPUT_LEN" \
|
||||
--ignore-eos \
|
||||
--disable-tqdm \
|
||||
--request-rate $request_rate \
|
||||
--percentile-metrics ttft,tpot,itl,e2el \
|
||||
--goodput e2el:$MAX_LATENCY_ALLOWED_MS \
|
||||
--goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \
|
||||
--num-prompts 100 \
|
||||
--random-prefix-len $prefix_len \
|
||||
--host "$HOSTNAME" \
|
||||
@@ -255,7 +255,7 @@ gpu_memory_utilization=0.98
|
||||
find_gpu_memory_utilization=0
|
||||
while (( $(echo "$gpu_memory_utilization >= 0.9" | bc -l) )); do
|
||||
# Pass empty string for profile_dir argument
|
||||
start_server $gpu_memory_utilization "${num_seqs_list[-1]}" "${num_batched_tokens_list[-1]}" "$LOG_FOLDER/vllm_log_gpu_memory_utilization_$gpu_memory_utilization.log" ""
|
||||
start_server "$gpu_memory_utilization" "${num_seqs_list[-1]}" "${num_batched_tokens_list[-1]}" "$LOG_FOLDER/vllm_log_gpu_memory_utilization_$gpu_memory_utilization.log" ""
|
||||
result=$?
|
||||
if [[ "$result" -eq 0 ]]; then
|
||||
find_gpu_memory_utilization=1
|
||||
@@ -274,7 +274,7 @@ fi
|
||||
|
||||
for num_seqs in "${num_seqs_list[@]}"; do
|
||||
for num_batched_tokens in "${num_batched_tokens_list[@]}"; do
|
||||
run_benchmark $num_seqs $num_batched_tokens $gpu_memory_utilization
|
||||
run_benchmark "$num_seqs" "$num_batched_tokens" "$gpu_memory_utilization"
|
||||
done
|
||||
done
|
||||
echo "finish permutations"
|
||||
@@ -285,7 +285,7 @@ echo "finish permutations"
|
||||
if (( $(echo "$best_throughput > 0" | bc -l) )); then
|
||||
echo
|
||||
echo "Benchmark tuning finished. Now running profiling on the best configuration found..."
|
||||
echo "Best config: max_num_seqs: $best_max_num_seqs, max_num_batched_tokens: $best_num_batched_tokens, throughput: $best_throughput"
|
||||
echo "Best config: max_num_seqs: $best_max_num_seqs, max_num_batched_tokens: $best_num_batched_tokens, throughput: $best_throughput, goodput: $best_goodput"
|
||||
echo
|
||||
|
||||
vllm_log="$LOG_FOLDER/vllm_log_BEST_PROFILE.txt"
|
||||
@@ -293,7 +293,7 @@ if (( $(echo "$best_throughput > 0" | bc -l) )); then
|
||||
|
||||
# Start server with the best params and profiling ENABLED
|
||||
echo "Starting server for profiling..."
|
||||
start_server $gpu_memory_utilization $best_max_num_seqs $best_num_batched_tokens "$vllm_log" "$PROFILE_PATH"
|
||||
start_server "$gpu_memory_utilization" "$best_max_num_seqs" "$best_num_batched_tokens" "$vllm_log" "$PROFILE_PATH"
|
||||
|
||||
# Run benchmark with the best params and the --profile flag
|
||||
echo "Running benchmark with profiling..."
|
||||
@@ -301,15 +301,15 @@ if (( $(echo "$best_throughput > 0" | bc -l) )); then
|
||||
adjusted_input_len=$(( INPUT_LEN - prefix_len ))
|
||||
vllm bench serve \
|
||||
--backend vllm \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--dataset-name random \
|
||||
--random-input-len $adjusted_input_len \
|
||||
--random-output-len $OUTPUT_LEN \
|
||||
--random-output-len "$OUTPUT_LEN" \
|
||||
--ignore-eos \
|
||||
--disable-tqdm \
|
||||
--request-rate $best_request_rate \
|
||||
--request-rate "$best_request_rate" \
|
||||
--percentile-metrics ttft,tpot,itl,e2el \
|
||||
--goodput e2el:$MAX_LATENCY_ALLOWED_MS \
|
||||
--goodput e2el:"$MAX_LATENCY_ALLOWED_MS" \
|
||||
--num-prompts 100 \
|
||||
--random-prefix-len $prefix_len \
|
||||
--host "$HOSTNAME" \
|
||||
|
||||
@@ -64,7 +64,7 @@ for i in $(seq 0 $(($num_runs - 1))); do
|
||||
else
|
||||
STATUS="FAILURE"
|
||||
((FAILURE_COUNT++))
|
||||
FAILED_RUNS+=("Run #$((i+1)): $(echo $run_object | jq -c .)")
|
||||
FAILED_RUNS+=("Run #$((i+1)): $(echo "$run_object" | jq -c .)")
|
||||
fi
|
||||
|
||||
RUN_OUTPUT=$(<"$RUN_OUTPUT_FILE")
|
||||
|
||||
@@ -0,0 +1,471 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Benchmark comparing Triton vs PyTorch sort-based top-k/top-p implementations.
|
||||
|
||||
Compares:
|
||||
- apply_top_k_top_p_triton (Triton binary search)
|
||||
- apply_top_k_top_p (PyTorch sort-based)
|
||||
|
||||
Scenarios:
|
||||
- top_k only (whole batch, partial batch)
|
||||
- top_p only (whole batch, partial batch)
|
||||
- mix of top_k and top_p
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch
|
||||
from vllm.v1.sample.ops.topk_topp_triton import (
|
||||
apply_top_k_top_p_triton,
|
||||
reset_buffer_cache,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkConfig:
|
||||
"""Configuration for a benchmark run."""
|
||||
|
||||
name: str
|
||||
batch_size: int
|
||||
vocab_size: int
|
||||
# k and p can be tensors or None
|
||||
k_values: torch.Tensor | None # [batch_size] or None
|
||||
p_values: torch.Tensor | None # [batch_size] or None
|
||||
description: str
|
||||
ops_pct: float = 0.0 # Percentage of ops relative to batch size
|
||||
|
||||
|
||||
def calculate_ops_pct(
|
||||
k_values: torch.Tensor | None,
|
||||
p_values: torch.Tensor | None,
|
||||
vocab_size: int,
|
||||
batch_size: int,
|
||||
) -> float:
|
||||
"""
|
||||
Calculate the percentage of active top-k and top-p operations.
|
||||
|
||||
Returns percentage where 100% = batch_size ops.
|
||||
E.g., if all rows have both top-k and top-p active, returns 200%.
|
||||
"""
|
||||
active_ops = 0
|
||||
|
||||
if k_values is not None:
|
||||
# Count rows where k < vocab_size (active top-k filtering)
|
||||
active_ops += (k_values < vocab_size).sum().item()
|
||||
|
||||
if p_values is not None:
|
||||
# Count rows where p < 1.0 (active top-p filtering)
|
||||
active_ops += (p_values < 1.0).sum().item()
|
||||
|
||||
return (active_ops / batch_size) * 100 if batch_size > 0 else 0.0
|
||||
|
||||
|
||||
def create_logits(
|
||||
batch_size: int, vocab_size: int, device: str = "cuda"
|
||||
) -> torch.Tensor:
|
||||
"""Create random logits mimicking a realistic LLM distribution.
|
||||
|
||||
Uses a Zipf-like probability distribution (rank^-1.1) converted to logits
|
||||
via log, then randomly permuted per row. This produces a peaked distribution
|
||||
where a small number of tokens capture most probability mass, similar to
|
||||
real model outputs.
|
||||
"""
|
||||
# Create Zipf-like probabilities: p(rank) ~ rank^(-alpha)
|
||||
ranks = torch.arange(1, vocab_size + 1, dtype=torch.float32, device=device)
|
||||
probs = ranks.pow(-1.1)
|
||||
probs = probs / probs.sum()
|
||||
|
||||
# Convert to logits (log-probabilities, unnormalized is fine)
|
||||
base_logits = probs.log()
|
||||
|
||||
# Broadcast to batch and randomly permute each row
|
||||
logits = base_logits.unsqueeze(0).expand(batch_size, -1).clone()
|
||||
for i in range(batch_size):
|
||||
logits[i] = logits[i, torch.randperm(vocab_size, device=device)]
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
def measure_memory() -> tuple[int, int]:
|
||||
"""Return (allocated, reserved) memory in bytes."""
|
||||
torch.cuda.synchronize()
|
||||
return torch.cuda.memory_allocated(), torch.cuda.max_memory_allocated()
|
||||
|
||||
|
||||
def reset_memory_stats():
|
||||
"""Reset peak memory statistics."""
|
||||
reset_buffer_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
|
||||
def benchmark_function(
|
||||
func,
|
||||
logits: torch.Tensor,
|
||||
k: torch.Tensor | None,
|
||||
p: torch.Tensor | None,
|
||||
warmup_iters: int = 5,
|
||||
benchmark_iters: int = 20,
|
||||
) -> tuple[float, int]:
|
||||
"""
|
||||
Benchmark a function and return (avg_time_ms, peak_memory_bytes).
|
||||
|
||||
Returns average time in milliseconds and peak memory usage.
|
||||
"""
|
||||
# Warmup
|
||||
for _ in range(warmup_iters):
|
||||
logits_copy = logits.clone()
|
||||
func(logits_copy, k, p)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Reset memory stats before benchmark
|
||||
reset_memory_stats()
|
||||
|
||||
# Benchmark
|
||||
start_events = [
|
||||
torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)
|
||||
]
|
||||
end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)]
|
||||
|
||||
for i in range(benchmark_iters):
|
||||
logits_copy = logits.clone()
|
||||
start_events[i].record()
|
||||
func(logits_copy, k, p)
|
||||
end_events[i].record()
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
# Calculate timing
|
||||
times = [
|
||||
start_events[i].elapsed_time(end_events[i]) for i in range(benchmark_iters)
|
||||
]
|
||||
avg_time = sum(times) / len(times)
|
||||
|
||||
# Get peak memory
|
||||
_, peak_memory = measure_memory()
|
||||
|
||||
return avg_time, peak_memory
|
||||
|
||||
|
||||
def create_benchmark_configs(
|
||||
batch_sizes: list[int],
|
||||
vocab_sizes: list[int],
|
||||
device: str = "cuda",
|
||||
) -> list[BenchmarkConfig]:
|
||||
"""Create all benchmark configurations."""
|
||||
configs = []
|
||||
|
||||
for vocab_size in vocab_sizes:
|
||||
for batch_size in batch_sizes:
|
||||
# 1. Top-k only - whole batch (all rows have k < vocab_size)
|
||||
k_all = torch.full((batch_size,), 50, dtype=torch.int32, device=device)
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"topk_whole_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=k_all,
|
||||
p_values=None,
|
||||
description=f"Top-k only (whole batch, k=50), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(k_all, None, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
# 2. Top-k only - partial batch (half have k=50, half have k=vocab_size)
|
||||
k_partial = torch.full((batch_size,), 50, dtype=torch.int32, device=device)
|
||||
k_partial[batch_size // 2 :] = vocab_size # No filtering for second half
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"topk_partial_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=k_partial,
|
||||
p_values=None,
|
||||
description=f"Top-k only (partial batch, 50% k=50, 50% k=vocab), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(k_partial, None, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
# 3. Top-p only - whole batch (all rows have p < 1.0)
|
||||
p_all = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device)
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"topp_whole_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=None,
|
||||
p_values=p_all,
|
||||
description=f"Top-p only (whole batch, p=0.9), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(None, p_all, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
# 4. Top-p only - partial batch (half have p=0.9, half have p=1.0)
|
||||
p_partial = torch.full(
|
||||
(batch_size,), 0.9, dtype=torch.float32, device=device
|
||||
)
|
||||
p_partial[batch_size // 2 :] = 1.0 # No filtering for second half
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"topp_partial_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=None,
|
||||
p_values=p_partial,
|
||||
description=f"Top-p only (partial batch, 50% p=0.9, 50% p=1.0), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(None, p_partial, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
# 5. Mix of top-k and top-p (both applied to whole batch)
|
||||
k_mix = torch.full((batch_size,), 100, dtype=torch.int32, device=device)
|
||||
p_mix = torch.full((batch_size,), 0.9, dtype=torch.float32, device=device)
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"topk_topp_whole_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=k_mix,
|
||||
p_values=p_mix,
|
||||
description=f"Top-k + Top-p (whole batch, k=100, p=0.9), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(k_mix, p_mix, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
# 6. Mix with partial application (some rows k only, some p only, some both)
|
||||
k_mixed = torch.full(
|
||||
(batch_size,), vocab_size, dtype=torch.int32, device=device
|
||||
)
|
||||
p_mixed = torch.full((batch_size,), 1.0, dtype=torch.float32, device=device)
|
||||
# First third: k only
|
||||
third = batch_size // 3
|
||||
k_mixed[:third] = 50
|
||||
# Second third: p only
|
||||
p_mixed[third : 2 * third] = 0.5
|
||||
# Last third: both k and p
|
||||
k_mixed[2 * third :] = 100
|
||||
p_mixed[2 * third :] = 0.9
|
||||
configs.append(
|
||||
BenchmarkConfig(
|
||||
name=f"mixed_partial_b{batch_size}_v{vocab_size // 1000}k",
|
||||
batch_size=batch_size,
|
||||
vocab_size=vocab_size,
|
||||
k_values=k_mixed,
|
||||
p_values=p_mixed,
|
||||
description=f"Mixed partial (1/3 k=50, 1/3 p=0.9, 1/3 both), "
|
||||
f"batch={batch_size}, vocab={vocab_size}",
|
||||
ops_pct=calculate_ops_pct(k_mixed, p_mixed, vocab_size, batch_size),
|
||||
)
|
||||
)
|
||||
|
||||
return configs
|
||||
|
||||
|
||||
def format_memory(bytes_val: int) -> str:
|
||||
"""Format memory in human-readable form."""
|
||||
if bytes_val >= 1024**3:
|
||||
return f"{bytes_val / (1024**3):.2f} GB"
|
||||
elif bytes_val >= 1024**2:
|
||||
return f"{bytes_val / (1024**2):.2f} MB"
|
||||
elif bytes_val >= 1024:
|
||||
return f"{bytes_val / 1024:.2f} KB"
|
||||
return f"{bytes_val} B"
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
configs: list[BenchmarkConfig],
|
||||
warmup_iters: int = 5,
|
||||
benchmark_iters: int = 20,
|
||||
verbose: bool = True,
|
||||
):
|
||||
"""Run all benchmarks and print results."""
|
||||
results = []
|
||||
|
||||
print("=" * 100)
|
||||
print("Top-k/Top-p Benchmark: Triton vs PyTorch Sort-based")
|
||||
print("=" * 100)
|
||||
print()
|
||||
|
||||
for config in configs:
|
||||
if verbose:
|
||||
print(f"Running: {config.description}")
|
||||
|
||||
# Create fresh logits for this config
|
||||
logits = create_logits(config.batch_size, config.vocab_size)
|
||||
|
||||
# Benchmark Triton
|
||||
reset_memory_stats()
|
||||
triton_time, triton_mem = benchmark_function(
|
||||
apply_top_k_top_p_triton,
|
||||
logits,
|
||||
config.k_values,
|
||||
config.p_values,
|
||||
warmup_iters,
|
||||
benchmark_iters,
|
||||
)
|
||||
|
||||
# Benchmark PyTorch
|
||||
reset_memory_stats()
|
||||
pytorch_time, pytorch_mem = benchmark_function(
|
||||
apply_top_k_top_p_pytorch,
|
||||
logits,
|
||||
config.k_values,
|
||||
config.p_values,
|
||||
warmup_iters,
|
||||
benchmark_iters,
|
||||
)
|
||||
|
||||
speedup = pytorch_time / triton_time if triton_time > 0 else float("inf")
|
||||
mem_ratio = pytorch_mem / triton_mem if triton_mem > 0 else float("inf")
|
||||
|
||||
result = {
|
||||
"config": config,
|
||||
"triton_time_ms": triton_time,
|
||||
"pytorch_time_ms": pytorch_time,
|
||||
"triton_mem": triton_mem,
|
||||
"pytorch_mem": pytorch_mem,
|
||||
"speedup": speedup,
|
||||
"mem_ratio": mem_ratio,
|
||||
}
|
||||
results.append(result)
|
||||
|
||||
if verbose:
|
||||
print(f" Triton: {triton_time:.3f} ms, {format_memory(triton_mem)}")
|
||||
print(f" PyTorch: {pytorch_time:.3f} ms, {format_memory(pytorch_mem)}")
|
||||
print(f" Speedup: {speedup:.2f}x, Memory ratio: {mem_ratio:.2f}x")
|
||||
print()
|
||||
|
||||
# Clean up
|
||||
del logits
|
||||
reset_memory_stats()
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def print_summary_table(results: list[dict]):
|
||||
"""Print a summary table of results."""
|
||||
print()
|
||||
print("=" * 130)
|
||||
print("SUMMARY TABLE")
|
||||
print("=" * 130)
|
||||
print()
|
||||
|
||||
# Header
|
||||
header = (
|
||||
f"{'Scenario':<40} {'Batch':>6} {'Vocab':>7} {'Ops%':>6} "
|
||||
f"{'Triton (ms)':>12} {'PyTorch (ms)':>13} {'Speedup':>8} "
|
||||
f"{'Tri Mem':>10} {'Pyt Mem':>10}"
|
||||
)
|
||||
print(header)
|
||||
print("-" * 130)
|
||||
|
||||
# Group by scenario type
|
||||
current_vocab = None
|
||||
for result in results:
|
||||
config = result["config"]
|
||||
|
||||
# Add separator between vocab sizes
|
||||
if current_vocab != config.vocab_size:
|
||||
if current_vocab is not None:
|
||||
print("-" * 130)
|
||||
current_vocab = config.vocab_size
|
||||
|
||||
scenario = config.name.split("_b")[0] # Extract scenario name
|
||||
print(
|
||||
f"{scenario:<40} {config.batch_size:>6} {config.vocab_size:>7} "
|
||||
f"{config.ops_pct:>5.0f}% "
|
||||
f"{result['triton_time_ms']:>12.3f} {result['pytorch_time_ms']:>13.3f} "
|
||||
f"{result['speedup']:>7.2f}x "
|
||||
f"{format_memory(result['triton_mem']):>10} "
|
||||
f"{format_memory(result['pytorch_mem']):>10}"
|
||||
)
|
||||
|
||||
print("=" * 130)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark Triton vs PyTorch sort-based top-k/top-p implementations"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-sizes",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[1, 4, 16, 64, 128, 512, 1024, 2048],
|
||||
help="Batch sizes to test (default: 1 4 16 64)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vocab-sizes",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[32768, 131072], # 32k, 128k
|
||||
help="Vocabulary sizes to test (default: 32768 131072)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-iters",
|
||||
type=int,
|
||||
default=5,
|
||||
help="Number of warmup iterations (default: 5)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--benchmark-iters",
|
||||
type=int,
|
||||
default=20,
|
||||
help="Number of benchmark iterations (default: 20)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quiet",
|
||||
action="store_true",
|
||||
help="Only print summary table",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Print configuration
|
||||
print(f"Batch sizes: {args.batch_sizes}")
|
||||
print(f"Vocab sizes: {args.vocab_sizes}")
|
||||
print(f"Warmup iterations: {args.warmup_iters}")
|
||||
print(f"Benchmark iterations: {args.benchmark_iters}")
|
||||
print()
|
||||
|
||||
# Check CUDA
|
||||
if not torch.cuda.is_available():
|
||||
print("ERROR: CUDA is not available. This benchmark requires a GPU.")
|
||||
return
|
||||
|
||||
device_name = torch.cuda.get_device_name(0)
|
||||
print(f"GPU: {device_name}")
|
||||
print()
|
||||
|
||||
# Create configs
|
||||
configs = create_benchmark_configs(
|
||||
args.batch_sizes,
|
||||
args.vocab_sizes,
|
||||
)
|
||||
|
||||
# Run benchmarks
|
||||
results = run_benchmark(
|
||||
configs,
|
||||
warmup_iters=args.warmup_iters,
|
||||
benchmark_iters=args.benchmark_iters,
|
||||
verbose=not args.quiet,
|
||||
)
|
||||
|
||||
# Print summary
|
||||
print_summary_table(results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -71,7 +71,7 @@ while [[ $# -gt 0 ]]; do
|
||||
usage
|
||||
;;
|
||||
*)
|
||||
echo "Unknown argument: $1\n"
|
||||
printf "Unknown argument: %s\n" "$1"
|
||||
usage
|
||||
;;
|
||||
esac
|
||||
@@ -84,15 +84,17 @@ mkdir -p "$OUTPUT_DIR"
|
||||
QPS_VALUES=(25 20 15 10 5 1)
|
||||
|
||||
# Common parameters
|
||||
COMMON_PARAMS="--backend $BACKEND \
|
||||
--model $MODEL \
|
||||
--dataset $DATASET \
|
||||
--structured-output-ratio $STRUCTURED_OUTPUT_RATIO \
|
||||
--save-results \
|
||||
--result-dir $OUTPUT_DIR \
|
||||
--output-len $MAX_NEW_TOKENS \
|
||||
--port $PORT \
|
||||
--tokenizer-mode $TOKENIZER_MODE"
|
||||
COMMON_PARAMS=(
|
||||
--backend "$BACKEND"
|
||||
--model "$MODEL"
|
||||
--dataset "$DATASET"
|
||||
--structured-output-ratio "$STRUCTURED_OUTPUT_RATIO"
|
||||
--save-results
|
||||
--result-dir "$OUTPUT_DIR"
|
||||
--output-len "$MAX_NEW_TOKENS"
|
||||
--port "$PORT"
|
||||
--tokenizer-mode "$TOKENIZER_MODE"
|
||||
)
|
||||
|
||||
echo "Starting structured output benchmark with model: $MODEL"
|
||||
echo "Backend: $BACKEND"
|
||||
@@ -109,17 +111,17 @@ for qps in "${QPS_VALUES[@]}"; do
|
||||
GIT_BRANCH=$(git rev-parse --abbrev-ref HEAD 2>/dev/null || echo "unknown")
|
||||
|
||||
# Construct filename for this run
|
||||
FILENAME="${BACKEND}_${qps}qps_$(basename $MODEL)_${DATASET}_${GIT_HASH}.json"
|
||||
FILENAME="${BACKEND}_${qps}qps_$(basename "$MODEL")_${DATASET}_${GIT_HASH}_${GIT_BRANCH}.json"
|
||||
|
||||
NUM_PROMPTS=$(echo "$TOTAL_SECONDS * $qps" | bc)
|
||||
NUM_PROMPTS=${NUM_PROMPTS%.*} # Remove fractional part
|
||||
echo "Running benchmark with $NUM_PROMPTS prompts"
|
||||
|
||||
# Run the benchmark
|
||||
python "$SCRIPT_DIR/benchmark_serving_structured_output.py" $COMMON_PARAMS \
|
||||
--request-rate $qps \
|
||||
python "$SCRIPT_DIR/benchmark_serving_structured_output.py" "${COMMON_PARAMS[@]}" \
|
||||
--request-rate "$qps" \
|
||||
--result-filename "$FILENAME" \
|
||||
--num-prompts $NUM_PROMPTS
|
||||
--num-prompts "$NUM_PROMPTS"
|
||||
|
||||
echo "Completed benchmark with QPS: $qps"
|
||||
echo "----------------------------------------"
|
||||
|
||||
@@ -18,6 +18,7 @@ set(ENABLE_AVX512 $ENV{VLLM_CPU_AVX512})
|
||||
set(ENABLE_AVX512BF16 $ENV{VLLM_CPU_AVX512BF16})
|
||||
set(ENABLE_AVX512VNNI $ENV{VLLM_CPU_AVX512VNNI})
|
||||
set(ENABLE_AMXBF16 $ENV{VLLM_CPU_AMXBF16})
|
||||
set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
|
||||
|
||||
include_directories("${CMAKE_SOURCE_DIR}/csrc")
|
||||
|
||||
@@ -115,6 +116,10 @@ else()
|
||||
set(AVX512_FOUND ON)
|
||||
message(STATUS "AVX512 support enabled via VLLM_CPU_AVX512 environment variable")
|
||||
endif()
|
||||
if (ENABLE_ARM_BF16)
|
||||
set(ARM_BF16_FOUND ON)
|
||||
message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (AVX512_FOUND AND NOT AVX512_DISABLED)
|
||||
|
||||
@@ -19,7 +19,7 @@ else()
|
||||
FetchContent_Declare(
|
||||
flashmla
|
||||
GIT_REPOSITORY https://github.com/vllm-project/FlashMLA
|
||||
GIT_TAG c2afa9cb93e674d5a9120a170a6da57b89267208
|
||||
GIT_TAG f46984b4caf19409bf4ab94a540341918dd6a1a3
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
|
||||
@@ -13,8 +13,13 @@ struct alignas(32) u32x8_t {
|
||||
uint32_t u0, u1, u2, u3, u4, u5, u6, u7;
|
||||
};
|
||||
|
||||
// 256-bit vector loads require both sm_100+ and CUDA 12.9+
|
||||
#define VLLM_SUPPORTS_256BIT_VECTORS \
|
||||
(defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 && \
|
||||
__CUDACC_VER_MAJOR__ >= 12 && __CUDACC_VER_MINOR__ >= 9)
|
||||
|
||||
__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
|
||||
#if VLLM_SUPPORTS_256BIT_VECTORS
|
||||
asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
|
||||
: "=r"(val.u0), "=r"(val.u1), "=r"(val.u2), "=r"(val.u3),
|
||||
"=r"(val.u4), "=r"(val.u5), "=r"(val.u6), "=r"(val.u7)
|
||||
@@ -35,7 +40,7 @@ __device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
|
||||
#if VLLM_SUPPORTS_256BIT_VECTORS
|
||||
asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
|
||||
:
|
||||
: "l"(ptr), "r"(val.u0), "r"(val.u1), "r"(val.u2), "r"(val.u3),
|
||||
|
||||
@@ -1305,7 +1305,8 @@ void indexer_k_quant_and_cache(
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(k));
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), "fp8_e4m3",
|
||||
static const std::string kv_cache_dtype = "fp8_e4m3";
|
||||
DISPATCH_BY_KV_CACHE_DTYPE(k.dtype(), kv_cache_dtype,
|
||||
CALL_INDEXER_K_QUANT_AND_CACHE);
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,747 @@
|
||||
/*
|
||||
* Adapted from
|
||||
* https://github.com/sgl-project/sglang/blob/main/sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu
|
||||
* which was adapted from
|
||||
* https://github.com/NVIDIA/TensorRT-LLM/blob/619709fc33bd5dc268f19d6a741fe7ed51c0f8f5/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3FusedAGemm.cu
|
||||
*
|
||||
* Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved.
|
||||
* Copyright (c) 2021, NAVER Corp. Authored by CLOVA.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include "core/registration.h"
|
||||
|
||||
#include <cstdlib>
|
||||
#include <mutex>
|
||||
|
||||
namespace {
|
||||
|
||||
inline int getSMVersion() {
|
||||
auto* props = at::cuda::getCurrentDeviceProperties();
|
||||
return props->major * 10 + props->minor;
|
||||
}
|
||||
|
||||
inline bool getEnvEnablePDL() {
|
||||
static std::once_flag flag;
|
||||
static bool enablePDL = false;
|
||||
std::call_once(flag, [&]() {
|
||||
if (getSMVersion() >= 90) {
|
||||
char const* env = std::getenv("TRTLLM_ENABLE_PDL");
|
||||
enablePDL = env && env[0] == '1' && env[1] == '\0';
|
||||
}
|
||||
});
|
||||
return enablePDL;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
using bf16_t = __nv_bfloat16;
|
||||
|
||||
__device__ void hmma_16_8_16_f32acc_bf16ab(float (&d_reg)[4],
|
||||
const bf16_t (&a_reg)[8],
|
||||
const bf16_t (&b_reg)[4],
|
||||
float const (&c_reg)[4]) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t a0 = *reinterpret_cast<uint32_t const*>(a_reg + 0);
|
||||
uint32_t a1 = *reinterpret_cast<uint32_t const*>(a_reg + 2);
|
||||
uint32_t a2 = *reinterpret_cast<uint32_t const*>(a_reg + 4);
|
||||
uint32_t a3 = *reinterpret_cast<uint32_t const*>(a_reg + 6);
|
||||
uint32_t b0 = *reinterpret_cast<uint32_t const*>(b_reg + 0);
|
||||
uint32_t b1 = *reinterpret_cast<uint32_t const*>(b_reg + 2);
|
||||
asm volatile(
|
||||
"mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
|
||||
"{%0, %1, %2, %3},"
|
||||
"{%4, %5, %6, %7},"
|
||||
"{%8, %9},"
|
||||
"{%10, %11, %12, %13};\n"
|
||||
: "=f"(d_reg[0]), "=f"(d_reg[1]), "=f"(d_reg[2]), "=f"(d_reg[3])
|
||||
: "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "f"(d_reg[0]),
|
||||
"f"(d_reg[1]), "f"(d_reg[2]), "f"(d_reg[3]));
|
||||
#endif
|
||||
}
|
||||
|
||||
extern "C" {
|
||||
__device__ uint32_t __nvvm_get_smem_pointer(void*);
|
||||
}
|
||||
|
||||
__device__ void ldgsts_128(void const* gPtr, void* sPtr, uint32_t pred) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
if (pred) {
|
||||
uint32_t smemPtrAsUint32 = __nvvm_get_smem_pointer(sPtr);
|
||||
asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2;\n" ::"r"(
|
||||
smemPtrAsUint32),
|
||||
"l"(gPtr), "n"(16));
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void ldsm_x4(void* smem_ptr, uint32_t* reg_ptr) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
asm volatile(
|
||||
"ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
|
||||
: "=r"(reg_ptr[0]), "=r"(reg_ptr[1]), "=r"(reg_ptr[2]), "=r"(reg_ptr[3])
|
||||
: "r"(__nvvm_get_smem_pointer(smem_ptr)));
|
||||
#endif
|
||||
}
|
||||
|
||||
template <class Type>
|
||||
__device__ int apply_swizzle_343_on_elem_row_col(int row_idx_, int col_idx_) {
|
||||
uint32_t row_idx = *reinterpret_cast<uint32_t*>(&row_idx_);
|
||||
uint32_t col_idx = *reinterpret_cast<uint32_t*>(&col_idx_);
|
||||
row_idx = row_idx % 8;
|
||||
row_idx = row_idx * (16 / sizeof(Type));
|
||||
col_idx = col_idx ^ row_idx;
|
||||
return *reinterpret_cast<int*>(&col_idx);
|
||||
}
|
||||
|
||||
__device__ void initialize_barrier(
|
||||
uint64_t* smem_barrier, // 64 bits user-manged barrier in smem
|
||||
int thread_count =
|
||||
1) // Thread count expected to arrive/wait on this barrier
|
||||
{
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
|
||||
asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smem_int_ptr),
|
||||
"r"(thread_count));
|
||||
#endif
|
||||
}
|
||||
|
||||
// Barrier wait
|
||||
__device__ void wait_barrier(
|
||||
uint64_t* smem_barrier, // 64 bits user-manged barrier in smem
|
||||
int phase_bit) // Current phase bit the barrier waiting to flip
|
||||
{
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .pred P1;\n"
|
||||
"LAB_WAIT:\n"
|
||||
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
|
||||
"@P1 bra DONE;\n"
|
||||
"bra LAB_WAIT;\n"
|
||||
"DONE:\n"
|
||||
"}\n" ::"r"(smem_int_ptr),
|
||||
"r"(phase_bit));
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ bool try_wait_barrier(uint64_t* smem_ptr, int phase_bit) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t wait_complete;
|
||||
uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_ptr);
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred P1; \n\t"
|
||||
"mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t"
|
||||
"selp.b32 %0, 1, 0, P1; \n\t"
|
||||
"}"
|
||||
: "=r"(wait_complete)
|
||||
: "r"(smem_int_ptr), "r"(phase_bit));
|
||||
return static_cast<bool>(wait_complete);
|
||||
#endif
|
||||
return false;
|
||||
}
|
||||
|
||||
// Barrier arrive
|
||||
__device__ void arrive_barrier(
|
||||
uint64_t* smem_barrier) // 64 bits user-manged barrier in smem
|
||||
{
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
|
||||
asm volatile(
|
||||
"{\n"
|
||||
".reg .b64 state; \n"
|
||||
"mbarrier.arrive.shared::cta.b64 state, [%0];\n"
|
||||
"}\n" ::"r"(smem_int_ptr));
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void ldgsts_arrive(uint64_t* smem_barrier) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier);
|
||||
asm volatile("cp.async.mbarrier.arrive.noinc.shared.b64 [%0];"
|
||||
:
|
||||
: "r"(smem_int_ptr));
|
||||
#endif
|
||||
}
|
||||
|
||||
template <int gemm_k, int tile_m, int tile_k, int stage_cnt>
|
||||
struct GmemLoaderA {
|
||||
static constexpr int elem_bytes = 2;
|
||||
static constexpr int vec_bytes = 16;
|
||||
static constexpr int vec_elems = vec_bytes / elem_bytes;
|
||||
static constexpr int thread_cnt = 64;
|
||||
static_assert((tile_m * tile_k) % (vec_elems * thread_cnt) == 0);
|
||||
static constexpr int a_inst_cnt_per_iter =
|
||||
(tile_m * tile_k) / (vec_elems * thread_cnt);
|
||||
static_assert(gemm_k % tile_k == 0);
|
||||
static constexpr int k_iter_cnt = gemm_k / tile_k;
|
||||
|
||||
// Extra params to keep the order of k reduction...
|
||||
static constexpr int mma_warp_cnt = 4;
|
||||
static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
|
||||
static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
|
||||
|
||||
private:
|
||||
__device__ int k_project(int tile_k_idx) {
|
||||
return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
|
||||
(tile_k_idx % per_mma_warp_k);
|
||||
}
|
||||
|
||||
public:
|
||||
__device__ GmemLoaderA(bf16_t const* gmem_a_local_, bf16_t* smem_a_,
|
||||
uint64_t* smem_barrier_)
|
||||
: gmem_a(gmem_a_local_),
|
||||
smem_a(smem_a_),
|
||||
smem_barrier(smem_barrier_),
|
||||
local_tid(threadIdx.x % thread_cnt) {}
|
||||
|
||||
__device__ void prepare() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
// swizzle, that's what we want.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < a_inst_cnt_per_iter; i++) {
|
||||
int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
|
||||
int m_idx = linear_idx / tile_k;
|
||||
int k_idx = linear_idx % tile_k;
|
||||
k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
|
||||
a_smem_offsets[i] = m_idx * tile_k + k_idx;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void issue_mainloop() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
#pragma unroll 1
|
||||
for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
|
||||
if (need_wait) {
|
||||
wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
|
||||
}
|
||||
int next_stage_idx = stage_idx + 1;
|
||||
int next_phase_bit =
|
||||
next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
|
||||
next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
|
||||
if (loop_idx != k_iter_cnt - 1) {
|
||||
need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
|
||||
next_phase_bit);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < a_inst_cnt_per_iter; i++) {
|
||||
int smem_offset = a_smem_offsets[i];
|
||||
bf16_t* smem_ptr_this_iter =
|
||||
smem_a + stage_idx * tile_m * tile_k + smem_offset;
|
||||
int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
|
||||
int m_idx = linear_idx / tile_k;
|
||||
int k_idx = linear_idx % tile_k;
|
||||
int gmem_offset = m_idx * gemm_k + k_project(k_idx);
|
||||
bf16_t const* gmem_ptr_this_iter = gmem_a + gmem_offset;
|
||||
ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, true);
|
||||
}
|
||||
ldgsts_arrive(smem_barrier + stage_idx * 2);
|
||||
|
||||
stage_idx = next_stage_idx;
|
||||
phase_bit = next_phase_bit;
|
||||
gmem_a += per_mma_warp_k;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
bf16_t const* gmem_a;
|
||||
bf16_t* smem_a;
|
||||
uint64_t* smem_barrier;
|
||||
int local_tid;
|
||||
int stage_idx = 0;
|
||||
int phase_bit = 1;
|
||||
bool need_wait = true;
|
||||
|
||||
// per smem_stage, store with swizzle information
|
||||
int a_smem_offsets[a_inst_cnt_per_iter];
|
||||
};
|
||||
|
||||
template <int gemm_k, int tile_n, int tile_k, int stage_cnt>
|
||||
struct GmemLoaderB {
|
||||
static constexpr int elem_bytes = 2;
|
||||
static constexpr int vec_bytes = 16;
|
||||
static constexpr int vec_elems = vec_bytes / elem_bytes;
|
||||
static constexpr int thread_cnt = 64;
|
||||
static_assert((tile_n * tile_k) % (vec_elems * thread_cnt) == 0);
|
||||
static constexpr int b_inst_cnt_per_iter =
|
||||
(tile_n * tile_k) / (vec_elems * thread_cnt);
|
||||
static_assert(gemm_k % tile_k == 0);
|
||||
static constexpr int k_iter_cnt = gemm_k / tile_k;
|
||||
|
||||
// Extra params to keep the order of k reduction...
|
||||
static constexpr int mma_warp_cnt = 4;
|
||||
static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt;
|
||||
static constexpr int k_each_chunk = gemm_k / mma_warp_cnt;
|
||||
|
||||
private:
|
||||
__device__ int k_project(int tile_k_idx) {
|
||||
return (tile_k_idx / per_mma_warp_k * k_each_chunk) +
|
||||
(tile_k_idx % per_mma_warp_k);
|
||||
}
|
||||
|
||||
public:
|
||||
__device__ GmemLoaderB(bf16_t const* gmem_b_local_, bf16_t* smem_b_,
|
||||
uint64_t* smem_barrier_, int gemm_n_)
|
||||
: gmem_b(gmem_b_local_),
|
||||
smem_b(smem_b_),
|
||||
smem_barrier(smem_barrier_),
|
||||
gemm_n(gemm_n_),
|
||||
local_tid(threadIdx.x % thread_cnt) {}
|
||||
|
||||
__device__ void prepare() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
// swizzle, that's what we want.
|
||||
#pragma unroll
|
||||
for (int i = 0; i < b_inst_cnt_per_iter; i++) {
|
||||
int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
|
||||
int n_idx = linear_idx / tile_k;
|
||||
int k_idx = linear_idx % tile_k;
|
||||
k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
|
||||
b_smem_offsets[i] = n_idx * tile_k + k_idx;
|
||||
preds[i] = n_idx < gemm_n;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void issue_mainloop() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
asm volatile("griddepcontrol.wait;");
|
||||
#pragma unroll 1
|
||||
for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
|
||||
if (need_wait) {
|
||||
wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit);
|
||||
}
|
||||
int next_stage_idx = stage_idx + 1;
|
||||
int next_phase_bit =
|
||||
next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
|
||||
next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx;
|
||||
if (loop_idx != k_iter_cnt - 1) {
|
||||
need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2,
|
||||
next_phase_bit);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < b_inst_cnt_per_iter; i++) {
|
||||
int smem_offset = b_smem_offsets[i];
|
||||
bf16_t* smem_ptr_this_iter =
|
||||
smem_b + stage_idx * tile_n * tile_k + smem_offset;
|
||||
int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems;
|
||||
int n_idx = linear_idx / tile_k;
|
||||
int k_idx = linear_idx % tile_k;
|
||||
int gmem_offset = n_idx * gemm_k + k_project(k_idx);
|
||||
bf16_t const* gmem_ptr_this_iter = gmem_b + gmem_offset;
|
||||
ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, preds[i]);
|
||||
}
|
||||
ldgsts_arrive(smem_barrier + stage_idx * 2);
|
||||
|
||||
stage_idx = next_stage_idx;
|
||||
phase_bit = next_phase_bit;
|
||||
gmem_b += per_mma_warp_k;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
bf16_t const* gmem_b;
|
||||
bf16_t* smem_b;
|
||||
uint64_t* smem_barrier;
|
||||
int gemm_n;
|
||||
int local_tid;
|
||||
int stage_idx = 0;
|
||||
int phase_bit = 1;
|
||||
bool need_wait = true;
|
||||
|
||||
// per smem_stage, store with swizzle information
|
||||
int b_smem_offsets[b_inst_cnt_per_iter];
|
||||
uint32_t preds[b_inst_cnt_per_iter];
|
||||
};
|
||||
|
||||
template <int gemm_m, int gemm_k, int tile_m, int tile_n, int tile_k,
|
||||
int stage_cnt>
|
||||
struct MmaComputer {
|
||||
static constexpr int elem_bytes = 2;
|
||||
static constexpr int thread_cnt = 128;
|
||||
static_assert(gemm_k % tile_k == 0);
|
||||
static_assert(tile_k % (thread_cnt / 32) == 0);
|
||||
static constexpr int per_warp_tile_k = tile_k / (thread_cnt / 32);
|
||||
static constexpr int k_iter_cnt = gemm_k / tile_k;
|
||||
static constexpr int k_phase_cnt = per_warp_tile_k / 16;
|
||||
static constexpr int m_iter_cnt = (tile_m + 15) / 16;
|
||||
static constexpr int n_iter_cnt =
|
||||
(tile_n + 7) /
|
||||
8; // Possible to have non-1 n_iter_cnt for ab_swap m16 case.
|
||||
static_assert(m_iter_cnt == 1);
|
||||
static_assert(n_iter_cnt == 1 || n_iter_cnt == 2);
|
||||
|
||||
__device__ MmaComputer(bf16_t* gmem_c_local_, bf16_t* smem_a_,
|
||||
bf16_t* smem_b_, uint64_t* smem_barrier_,
|
||||
int warp_idx_, int gemm_n_)
|
||||
: gmem_c(gmem_c_local_),
|
||||
smem_a(smem_a_),
|
||||
smem_b(smem_b_),
|
||||
smem_barrier(smem_barrier_),
|
||||
warp_idx(warp_idx_ - (thread_cnt / 32)),
|
||||
gemm_n(gemm_n_) {}
|
||||
|
||||
private:
|
||||
__device__ constexpr int internal_b_atom_func(int tid) {
|
||||
if constexpr (tile_n < 8) {
|
||||
return (tid % tile_n) + ((tid % 8) / tile_n * 0) + tid / 8 * 8 * tile_n;
|
||||
} else {
|
||||
return (tid % 8) + ((tid % 32) / 8 * (tile_n * 8));
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
__device__ void prepare() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
#pragma unroll
|
||||
for (int i = 0; i < k_phase_cnt; i++) {
|
||||
int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256;
|
||||
int m_idx = linear_idx % tile_m;
|
||||
int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k;
|
||||
k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(m_idx, k_idx);
|
||||
a_smem_offsets[0][i] = m_idx * tile_k + k_idx;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < k_phase_cnt; i += 2) { // Special i+=2 for B.
|
||||
int linear_idx =
|
||||
internal_b_atom_func(lane_idx) + i * tile_n * 16 + n_iter_idx * 8;
|
||||
int n_idx = linear_idx % tile_n;
|
||||
int k_idx = linear_idx / tile_n + warp_k_offset_in_tile_k;
|
||||
k_idx = apply_swizzle_343_on_elem_row_col<bf16_t>(n_idx, k_idx);
|
||||
b_smem_offsets[n_iter_idx][i] = n_idx * tile_k + k_idx;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void issue_mainloop() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
#pragma unroll 1
|
||||
for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) {
|
||||
wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit);
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0; i < k_phase_cnt; i++) {
|
||||
int smem_offset = a_smem_offsets[0][i];
|
||||
bf16_t* smem_ptr_this_iter =
|
||||
smem_a + stage_idx * tile_m * tile_k + smem_offset;
|
||||
ldsm_x4(smem_ptr_this_iter, reinterpret_cast<uint32_t*>(a_reg[0][i]));
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < k_phase_cnt; i += 2) {
|
||||
int smem_offset = b_smem_offsets[n_iter_idx][i];
|
||||
bf16_t* smem_ptr_this_iter =
|
||||
smem_b + stage_idx * tile_n * tile_k + smem_offset;
|
||||
ldsm_x4(smem_ptr_this_iter,
|
||||
reinterpret_cast<uint32_t*>(b_reg[n_iter_idx][i]));
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) {
|
||||
#pragma unroll
|
||||
for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) {
|
||||
hmma_16_8_16_f32acc_bf16ab(
|
||||
acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx],
|
||||
b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]);
|
||||
}
|
||||
}
|
||||
::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
|
||||
stage_idx += 1;
|
||||
phase_bit = stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
|
||||
stage_idx = stage_idx == stage_cnt ? 0 : stage_idx;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ void epi() {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
|
||||
// reorganize the acc_reg
|
||||
constexpr int thread_m = 2;
|
||||
constexpr int thread_n = 2 * n_iter_cnt;
|
||||
constexpr int cta_mma_n = n_iter_cnt * 8;
|
||||
float acc_reg_reorg[thread_m][thread_n];
|
||||
|
||||
for (int i = 0; i < thread_m; i++) {
|
||||
for (int j = 0; j < thread_n; j++) {
|
||||
acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)];
|
||||
}
|
||||
}
|
||||
|
||||
// 4 x cosize(smem_c_layout)
|
||||
float* smem_c = reinterpret_cast<float*>(smem_a);
|
||||
// coord -> index
|
||||
auto smem_c_index_func = [&](int m_idx, int n_idx) {
|
||||
int group_rows = 32 / cta_mma_n;
|
||||
int group_cnt = 2;
|
||||
return (m_idx % group_rows * cta_mma_n) +
|
||||
(m_idx / group_rows * (32 + group_cnt)) + n_idx;
|
||||
};
|
||||
constexpr int cosize_smem_c = ((tile_m * cta_mma_n) / 32) * (32 + 2);
|
||||
|
||||
// This should be optimized to STS.64 but can not be STS.128 due to the bank
|
||||
// index.
|
||||
#pragma unroll
|
||||
for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) {
|
||||
#pragma unroll
|
||||
for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) {
|
||||
int m_idx = (lane_idx / 4) + m_idx_thread * 8;
|
||||
int n_idx =
|
||||
((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8;
|
||||
smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] =
|
||||
acc_reg_reorg[m_idx_thread][n_idx_thread];
|
||||
}
|
||||
}
|
||||
asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt));
|
||||
|
||||
if (warp_idx == 0) {
|
||||
constexpr int final_acc_reg_cnt = (tile_m * tile_n + 31) / 32;
|
||||
float acc_final[final_acc_reg_cnt]{};
|
||||
|
||||
#pragma unroll
|
||||
for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
|
||||
int linear_idx = reg_idx * 32 + lane_idx;
|
||||
int m_idx = linear_idx % tile_m;
|
||||
int n_idx = linear_idx / tile_m;
|
||||
acc_final[reg_idx] +=
|
||||
smem_c[smem_c_index_func(m_idx, n_idx) + 0 * cosize_smem_c] +
|
||||
smem_c[smem_c_index_func(m_idx, n_idx) + 1 * cosize_smem_c] +
|
||||
smem_c[smem_c_index_func(m_idx, n_idx) + 2 * cosize_smem_c] +
|
||||
smem_c[smem_c_index_func(m_idx, n_idx) + 3 * cosize_smem_c];
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) {
|
||||
int linear_idx = reg_idx * 32 + lane_idx;
|
||||
int m_idx = linear_idx % tile_m;
|
||||
int n_idx = linear_idx / tile_m;
|
||||
if (m_idx < tile_m && n_idx < gemm_n) {
|
||||
gmem_c[n_idx * gemm_m + m_idx] = acc_final[reg_idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
bf16_t* gmem_c;
|
||||
bf16_t* smem_a;
|
||||
bf16_t* smem_b;
|
||||
uint64_t* smem_barrier;
|
||||
int warp_idx;
|
||||
int gemm_n;
|
||||
int stage_idx = 0;
|
||||
int phase_bit = 0;
|
||||
int lane_idx = threadIdx.x % 32;
|
||||
int warp_k_offset_in_tile_k = warp_idx * per_warp_tile_k;
|
||||
|
||||
int a_smem_offsets[m_iter_cnt][k_phase_cnt];
|
||||
int b_smem_offsets[n_iter_cnt][k_phase_cnt];
|
||||
|
||||
bf16_t a_reg[m_iter_cnt][k_phase_cnt][8];
|
||||
bf16_t b_reg[n_iter_cnt][k_phase_cnt][4];
|
||||
float acc_reg[m_iter_cnt][n_iter_cnt][4]{};
|
||||
};
|
||||
|
||||
// AB swapped, kernel is k-major, k-major, m-major
|
||||
template <int batch_size, int gemm_m, int gemm_k, int tile_m, int tile_n,
|
||||
int tile_k, int stage_cnt>
|
||||
__global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel(
|
||||
bf16_t* output, bf16_t const* mat_a, bf16_t const* mat_b, int gemm_n) {
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
|
||||
constexpr int load_thread_cnt = 128;
|
||||
constexpr int compute_thread_cnt = 128;
|
||||
constexpr int thread_cnt = load_thread_cnt + compute_thread_cnt;
|
||||
(void)thread_cnt;
|
||||
static_assert(gemm_m % 16 == 0);
|
||||
static_assert(gemm_k % tile_k == 0);
|
||||
static_assert(gemm_m % tile_m == 0);
|
||||
static_assert(
|
||||
tile_k == 128 || tile_k == 256 || tile_k == 512 ||
|
||||
tile_k == 1024); // tile_k must be larger than 64 since 4 warp splitK.
|
||||
static_assert(tile_m == 16);
|
||||
constexpr int g2s_vec_bytes = 16;
|
||||
constexpr int a_elem_bytes = 2;
|
||||
constexpr int b_elem_bytes = 2;
|
||||
static_assert((tile_m * a_elem_bytes + tile_n * b_elem_bytes) * tile_k *
|
||||
stage_cnt <=
|
||||
225 * 1024);
|
||||
static_assert((tile_m * tile_k * a_elem_bytes) %
|
||||
(load_thread_cnt * g2s_vec_bytes) ==
|
||||
0);
|
||||
static_assert((tile_n * tile_k * b_elem_bytes) %
|
||||
(load_thread_cnt * g2s_vec_bytes) ==
|
||||
0);
|
||||
|
||||
extern __shared__ char smem[];
|
||||
uint64_t* smem_barrier = reinterpret_cast<uint64_t*>(
|
||||
smem); // producer,consumer; producer,consumer; ...
|
||||
bf16_t* smem_a = reinterpret_cast<bf16_t*>(smem + (stage_cnt * 8 * 2 + 1024) /
|
||||
1024 * 1024);
|
||||
bf16_t* smem_b = smem_a + tile_m * tile_k * stage_cnt;
|
||||
|
||||
int cta_m_idx = tile_m * blockIdx.x;
|
||||
int cta_n_idx = tile_n * blockIdx.y;
|
||||
bf16_t const* gmem_a_local = mat_a + cta_m_idx * gemm_k;
|
||||
bf16_t const* gmem_b_local = mat_b + cta_n_idx * gemm_k;
|
||||
bf16_t* gmem_c_local = output + cta_n_idx * gemm_m + cta_m_idx;
|
||||
|
||||
int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
|
||||
|
||||
if (warp_idx == 4) {
|
||||
for (int i = 0; i < stage_cnt; i++) {
|
||||
initialize_barrier(smem_barrier + i * 2 + 0,
|
||||
load_thread_cnt); // producer
|
||||
initialize_barrier(smem_barrier + i * 2 + 1,
|
||||
compute_thread_cnt); // consumer
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (warp_idx < 2) {
|
||||
GmemLoaderA<gemm_k, tile_m, tile_k, stage_cnt> a_loader(
|
||||
gmem_a_local, smem_a, smem_barrier);
|
||||
a_loader.prepare();
|
||||
a_loader.issue_mainloop();
|
||||
} else if (warp_idx < 4) {
|
||||
GmemLoaderB<gemm_k, tile_n, tile_k, stage_cnt> b_loader(
|
||||
gmem_b_local, smem_b, smem_barrier, gemm_n);
|
||||
b_loader.prepare();
|
||||
b_loader.issue_mainloop();
|
||||
} else {
|
||||
MmaComputer<gemm_m, gemm_k, tile_m, tile_n, tile_k, stage_cnt> mma_computer(
|
||||
gmem_c_local, smem_a, smem_b, smem_barrier, warp_idx, gemm_n);
|
||||
mma_computer.prepare();
|
||||
mma_computer.issue_mainloop();
|
||||
mma_computer.epi();
|
||||
}
|
||||
asm volatile("griddepcontrol.launch_dependents;");
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename T, int kHdIn, int kHdOut, int kTileN>
|
||||
void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens,
|
||||
cudaStream_t const stream) {
|
||||
constexpr int gemm_m = kHdOut; // 2112
|
||||
int const gemm_n = num_tokens; // 1-16
|
||||
constexpr int gemm_k = kHdIn; // 7168
|
||||
constexpr int batch_size = 1;
|
||||
std::swap(mat_a, mat_b);
|
||||
constexpr int tile_m = 16;
|
||||
constexpr int tile_n = kTileN; // 8 or 16
|
||||
constexpr int tile_k = std::max(256, 1024 / tile_n); // 256
|
||||
constexpr int max_stage_cnt =
|
||||
1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t));
|
||||
constexpr int k_iter_cnt = gemm_k / tile_k;
|
||||
constexpr int stage_cnt =
|
||||
k_iter_cnt > max_stage_cnt ? max_stage_cnt : k_iter_cnt;
|
||||
int cta_m_cnt = gemm_m / tile_m;
|
||||
int cta_n_cnt = (gemm_n + tile_n - 1) / tile_n;
|
||||
constexpr int barrier_bytes = (stage_cnt * 16 + 1023) / 1024 * 1024;
|
||||
constexpr int smem_bytes =
|
||||
((tile_m * 2 + tile_n * 2) * tile_k * stage_cnt + barrier_bytes + 1023) /
|
||||
1024 * 1024;
|
||||
|
||||
dim3 grid(cta_m_cnt, cta_n_cnt, 1);
|
||||
dim3 block_size(256);
|
||||
cudaLaunchConfig_t config;
|
||||
config.gridDim = grid;
|
||||
config.blockDim = block_size;
|
||||
config.dynamicSmemBytes = smem_bytes;
|
||||
config.stream = stream;
|
||||
cudaLaunchAttribute attrs[1];
|
||||
attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
|
||||
attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL();
|
||||
config.numAttrs = 1;
|
||||
config.attrs = attrs;
|
||||
if (smem_bytes >= (48 * 1024)) {
|
||||
cudaFuncSetAttribute(fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
|
||||
tile_n, tile_k, stage_cnt>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
smem_bytes);
|
||||
}
|
||||
cudaLaunchKernelEx(&config,
|
||||
fused_a_gemm_kernel<batch_size, gemm_m, gemm_k, tile_m,
|
||||
tile_n, tile_k, stage_cnt>,
|
||||
output, mat_a, mat_b, gemm_n);
|
||||
}
|
||||
|
||||
template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>(
|
||||
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
|
||||
cudaStream_t);
|
||||
|
||||
template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>(
|
||||
__nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens,
|
||||
cudaStream_t);
|
||||
|
||||
void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
|
||||
torch::Tensor const& mat_b) {
|
||||
TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2);
|
||||
int const num_tokens = mat_a.size(0);
|
||||
int const hd_in = mat_a.size(1);
|
||||
int const hd_out = mat_b.size(1);
|
||||
|
||||
constexpr int kHdIn = 7168;
|
||||
constexpr int kHdOut = 2112;
|
||||
TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16,
|
||||
"required 1 <= mat_a.shape[0] <= 16")
|
||||
TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168")
|
||||
TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112")
|
||||
TORCH_CHECK(output.size(0) == num_tokens,
|
||||
"required output.shape[0] == mat_a.shape[0]")
|
||||
TORCH_CHECK(output.size(1) == hd_out,
|
||||
"required output.shape[1] == mat_b.shape[1]")
|
||||
|
||||
TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor");
|
||||
TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor");
|
||||
TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor");
|
||||
|
||||
TORCH_CHECK(mat_a.scalar_type() == torch::kBFloat16 &&
|
||||
mat_b.scalar_type() == torch::kBFloat16,
|
||||
"Only BFloat16 input dtype is supported")
|
||||
TORCH_CHECK(output.scalar_type() == torch::kBFloat16,
|
||||
"Only BFloat16 output dtype is supported")
|
||||
|
||||
TORCH_CHECK(getSMVersion() >= 90, "required CUDA ARCH >= SM_90");
|
||||
|
||||
auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device());
|
||||
if (num_tokens <= 8) {
|
||||
invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>(
|
||||
reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
|
||||
stream);
|
||||
} else {
|
||||
invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>(
|
||||
reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()),
|
||||
reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), num_tokens,
|
||||
stream);
|
||||
}
|
||||
}
|
||||
+8
-1
@@ -315,7 +315,9 @@ void silu_and_mul_scaled_fp4_experts_quant(
|
||||
void per_token_group_quant_fp8(const torch::Tensor& input,
|
||||
torch::Tensor& output_q, torch::Tensor& output_s,
|
||||
int64_t group_size, double eps, double fp8_min,
|
||||
double fp8_max, bool scale_ue8m0);
|
||||
double fp8_max, bool scale_ue8m0,
|
||||
bool dummy_is_scale_transposed,
|
||||
bool dummy_is_tma_aligned);
|
||||
|
||||
void per_token_group_quant_int8(const torch::Tensor& input,
|
||||
torch::Tensor& output_q,
|
||||
@@ -408,3 +410,8 @@ void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
|
||||
int64_t quant_level, bool cast_bf2half = false);
|
||||
int64_t qr_max_size();
|
||||
#endif
|
||||
|
||||
#ifndef USE_ROCM
|
||||
void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a,
|
||||
torch::Tensor const& mat_b);
|
||||
#endif
|
||||
@@ -97,7 +97,7 @@ __global__ void rms_norm_per_block_quant_kernel(
|
||||
scalar_t const* __restrict__ input, // [..., hidden_size]
|
||||
scalar_t const* __restrict__ weight, // [hidden_size]
|
||||
float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
|
||||
scalar_t* __restrict__ residual = nullptr) {
|
||||
scalar_t* __restrict__ residual = nullptr, int64_t outer_scale_stride = 1) {
|
||||
float rms;
|
||||
// Compute RMS
|
||||
// Always able to vectorize due to constraints on hidden_size
|
||||
@@ -108,7 +108,8 @@ __global__ void rms_norm_per_block_quant_kernel(
|
||||
// Always able to vectorize due to constraints on hidden_size and group_size
|
||||
vllm::vectorized::compute_dynamic_per_token_scales<
|
||||
scalar_t, scalar_out_t, has_residual, is_scale_transposed, group_size>(
|
||||
nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual);
|
||||
nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual,
|
||||
outer_scale_stride);
|
||||
|
||||
// RMS Norm + Quant
|
||||
// Always able to vectorize due to constraints on hidden_size
|
||||
@@ -119,7 +120,8 @@ __global__ void rms_norm_per_block_quant_kernel(
|
||||
vllm::vectorized::norm_and_quant<
|
||||
scalar_t, scalar_out_t, std::is_same_v<scalar_out_t, int8_t>,
|
||||
has_residual, is_scale_transposed, group_size>(
|
||||
out, input, weight, rms, scales, hidden_size, residual);
|
||||
out, input, weight, rms, scales, hidden_size, residual,
|
||||
outer_scale_stride);
|
||||
}
|
||||
|
||||
} // namespace vllm
|
||||
@@ -225,7 +227,8 @@ void rms_norm_per_block_quant_dispatch(
|
||||
: nullptr,
|
||||
var_epsilon, hidden_size,
|
||||
has_residual ? residual->data_ptr<scalar_in_t>()
|
||||
: nullptr);
|
||||
: nullptr,
|
||||
scales.stride(1));
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -257,6 +260,11 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
|
||||
TORCH_CHECK(group_size == 128 || group_size == 64,
|
||||
"Unsupported group size: ", group_size);
|
||||
|
||||
if (scales.stride(1) > 1) {
|
||||
TORCH_CHECK(is_scale_transposed,
|
||||
"Outer scale stride must be 1 when scales are not transposed");
|
||||
}
|
||||
|
||||
rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
|
||||
var_epsilon, scale_ub, residual,
|
||||
is_scale_transposed);
|
||||
|
||||
@@ -74,7 +74,7 @@ __device__ void compute_dynamic_per_token_scales(
|
||||
scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
|
||||
float const rms, float const* __restrict__ scale_ub,
|
||||
int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
|
||||
int32_t const group_size = 0) {
|
||||
int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
|
||||
float block_absmax_val_maybe = 0.0f;
|
||||
constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
|
||||
__syncthreads();
|
||||
@@ -133,7 +133,9 @@ __device__ void compute_dynamic_per_token_scales(
|
||||
scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
|
||||
// Global output store
|
||||
if constexpr (is_scale_transposed) {
|
||||
all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +
|
||||
int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
|
||||
outer_scale_stride * outer_scale_stride;
|
||||
all_token_scales[(threadIdx.x / threads_per_group) * scale_rows +
|
||||
blockIdx.x] = scale;
|
||||
} else {
|
||||
all_token_scales[blockIdx.x * num_groups +
|
||||
@@ -180,13 +182,11 @@ __device__ void compute_dynamic_per_token_scales(
|
||||
|
||||
template <typename scalar_t, typename scalar_out_t, bool is_scale_inverted,
|
||||
bool has_residual = false, bool is_scale_transposed = false>
|
||||
__device__ void norm_and_quant(scalar_out_t* __restrict__ output,
|
||||
scalar_t const* __restrict__ input,
|
||||
scalar_t const* __restrict__ weight,
|
||||
float const rms, float* const scale,
|
||||
int32_t const hidden_size,
|
||||
scalar_t* __restrict__ residual = nullptr,
|
||||
int32_t const group_size = 0) {
|
||||
__device__ void norm_and_quant(
|
||||
scalar_out_t* __restrict__ output, scalar_t const* __restrict__ input,
|
||||
scalar_t const* __restrict__ weight, float const rms, float* const scale,
|
||||
int32_t const hidden_size, scalar_t* __restrict__ residual = nullptr,
|
||||
int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
|
||||
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
|
||||
|
||||
for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
|
||||
@@ -202,7 +202,9 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
|
||||
int64_t scale_idx = 0;
|
||||
if (group_size > 0) {
|
||||
if constexpr (is_scale_transposed) {
|
||||
scale_idx = (i / group_size) * gridDim.x + blockIdx.x;
|
||||
int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
|
||||
outer_scale_stride * outer_scale_stride;
|
||||
scale_idx = (i / group_size) * scale_rows + blockIdx.x;
|
||||
} else {
|
||||
scale_idx = blockIdx.x * (hidden_size / group_size) + i / group_size;
|
||||
}
|
||||
@@ -286,8 +288,8 @@ __device__ void compute_dynamic_per_token_scales(
|
||||
float* __restrict__ token_scale, float* __restrict__ all_token_scales,
|
||||
scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
|
||||
float const rms, float const* __restrict__ scale_ub,
|
||||
int32_t const hidden_size,
|
||||
scalar_t const* __restrict__ residual = nullptr) {
|
||||
int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
|
||||
int64_t outer_scale_stride = 1) {
|
||||
constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
|
||||
|
||||
const int VEC_SIZE = 4;
|
||||
@@ -382,7 +384,9 @@ __device__ void compute_dynamic_per_token_scales(
|
||||
scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
|
||||
// Global output store
|
||||
if constexpr (is_scale_transposed) {
|
||||
all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +
|
||||
int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
|
||||
outer_scale_stride * outer_scale_stride;
|
||||
all_token_scales[(threadIdx.x / threads_per_group) * scale_rows +
|
||||
blockIdx.x] = scale;
|
||||
} else {
|
||||
all_token_scales[blockIdx.x * num_groups +
|
||||
@@ -463,7 +467,8 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
|
||||
scalar_t const* __restrict__ weight,
|
||||
float const rms, float* const scale,
|
||||
int32_t const hidden_size,
|
||||
scalar_t* __restrict__ residual = nullptr) {
|
||||
scalar_t* __restrict__ residual = nullptr,
|
||||
int64_t outer_scale_stride = 1) {
|
||||
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
|
||||
|
||||
// Vectorized input/output/weight/residual to better utilize memory bandwidth.
|
||||
@@ -516,7 +521,9 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
|
||||
int64_t const num_groups = hidden_size / group_size;
|
||||
int64_t scale_idx = 0;
|
||||
if constexpr (is_scale_transposed) {
|
||||
scale_idx = (i * VEC_SIZE / group_size) * gridDim.x + blockIdx.x;
|
||||
int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /
|
||||
outer_scale_stride * outer_scale_stride;
|
||||
scale_idx = (i * VEC_SIZE / group_size) * scale_rows + blockIdx.x;
|
||||
} else {
|
||||
scale_idx = blockIdx.x * num_groups + i * VEC_SIZE / group_size;
|
||||
}
|
||||
|
||||
@@ -379,7 +379,9 @@ void per_token_group_quant_8bit_packed(const torch::Tensor& input,
|
||||
void per_token_group_quant_fp8(const torch::Tensor& input,
|
||||
torch::Tensor& output_q, torch::Tensor& output_s,
|
||||
int64_t group_size, double eps, double fp8_min,
|
||||
double fp8_max, bool scale_ue8m0) {
|
||||
double fp8_max, bool scale_ue8m0,
|
||||
bool dummy_is_scale_transposed = false,
|
||||
bool dummy_is_tma_aligned = false) {
|
||||
per_token_group_quant_8bit(input, output_q, output_s, group_size, eps,
|
||||
fp8_min, fp8_max, scale_ue8m0);
|
||||
}
|
||||
@@ -239,6 +239,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
|
||||
// Quantization ops
|
||||
#ifndef USE_ROCM
|
||||
// DeepSeek V3 fused A GEMM (SM 9.0+, bf16 only, 1-16 tokens).
|
||||
ops.def(
|
||||
"dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()");
|
||||
ops.impl("dsv3_fused_a_gemm", torch::kCUDA, &dsv3_fused_a_gemm);
|
||||
|
||||
// Quantized GEMM for AWQ.
|
||||
ops.def(
|
||||
"awq_gemm(Tensor _in_feats, Tensor _kernel, Tensor _scaling_factors, "
|
||||
@@ -643,11 +648,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
|
||||
#ifndef USE_ROCM
|
||||
// Compute per-token-group FP8 quantized tensor and scaling factor.
|
||||
// The dummy arguments are here so we can correctly fuse with RMSNorm.
|
||||
ops.def(
|
||||
"per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! "
|
||||
"output_s, "
|
||||
"int group_size, float eps, float fp8_min, float fp8_max, bool "
|
||||
"scale_ue8m0) -> ()");
|
||||
"scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned "
|
||||
") -> ()");
|
||||
ops.impl("per_token_group_fp8_quant", torch::kCUDA,
|
||||
&per_token_group_quant_fp8);
|
||||
|
||||
|
||||
+4
-8
@@ -418,10 +418,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
fi && \
|
||||
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38
|
||||
|
||||
# Copy extension wheels from extensions-build stage for later use
|
||||
COPY --from=extensions-build /tmp/deepgemm/dist /tmp/deepgemm/dist
|
||||
COPY --from=extensions-build /tmp/ep_kernels_workspace/dist /tmp/ep_kernels_workspace/dist
|
||||
|
||||
# Check the size of the wheel if RUN_WHEEL_CHECK is true
|
||||
COPY .buildkite/check-wheel-size.py check-wheel-size.py
|
||||
# sync the default value with .buildkite/check-wheel-size.py
|
||||
@@ -660,9 +656,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
. /etc/environment && \
|
||||
uv pip list
|
||||
|
||||
# Install deepgemm wheel that has been built in the `build` stage
|
||||
# Install deepgemm wheel that has been built in the `extensions-build` stage
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
--mount=type=bind,from=build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \
|
||||
--mount=type=bind,from=extensions-build,source=/tmp/deepgemm/dist,target=/tmp/deepgemm/dist,ro \
|
||||
sh -c 'if ls /tmp/deepgemm/dist/*.whl >/dev/null 2>&1; then \
|
||||
uv pip install --system /tmp/deepgemm/dist/*.whl; \
|
||||
else \
|
||||
@@ -672,8 +668,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# Pytorch now installs NVSHMEM, setting LD_LIBRARY_PATH
|
||||
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
|
||||
|
||||
# Install EP kernels wheels (pplx-kernels and DeepEP) that have been built in the `build` stage
|
||||
RUN --mount=type=bind,from=build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
|
||||
# Install EP kernels wheels (pplx-kernels and DeepEP) that have been built in the `extensions-build` stage
|
||||
RUN --mount=type=bind,from=extensions-build,src=/tmp/ep_kernels_workspace/dist,target=/vllm-workspace/ep_kernels/dist \
|
||||
--mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --system ep_kernels/dist/*.whl --verbose \
|
||||
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
|
||||
|
||||
@@ -20,6 +20,7 @@
|
||||
# VLLM_CPU_AVX512BF16=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AVX512VNNI=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_AMXBF16=false (default)|true (for cross-compilation)
|
||||
# VLLM_CPU_ARM_BF16=false (default)|true (for cross-compilation)
|
||||
#
|
||||
|
||||
######################### COMMON BASE IMAGE #########################
|
||||
@@ -108,9 +109,22 @@ ENV VLLM_CPU_AVX512VNNI=${VLLM_CPU_AVX512VNNI}
|
||||
# Support for building with AMXBF16 ISA: docker build --build-arg VLLM_CPU_AMXBF16="true" ...
|
||||
ARG VLLM_CPU_AMXBF16=1
|
||||
ENV VLLM_CPU_AMXBF16=${VLLM_CPU_AMXBF16}
|
||||
# Support for cross-compilation with ARM BF16 ISA: docker build --build-arg VLLM_CPU_ARM_BF16="true" ...
|
||||
ARG VLLM_CPU_ARM_BF16=0
|
||||
ENV VLLM_CPU_ARM_BF16=${VLLM_CPU_ARM_BF16}
|
||||
|
||||
WORKDIR /vllm-workspace
|
||||
|
||||
# Validate build arguments - prevent mixing incompatible ISA flags
|
||||
RUN if [ "$TARGETARCH" = "arm64" ] && { [ "$VLLM_CPU_AVX2" != "0" ] || [ "$VLLM_CPU_AVX512" != "0" ] || [ "$VLLM_CPU_AVX512BF16" != "0" ] || [ "$VLLM_CPU_AVX512VNNI" != "0" ]; }; then \
|
||||
echo "ERROR: Cannot use x86-specific ISA flags (AVX2, AVX512, etc.) when building for ARM64 (--platform=linux/arm64)"; \
|
||||
exit 1; \
|
||||
fi && \
|
||||
if [ "$TARGETARCH" = "amd64" ] && [ "$VLLM_CPU_ARM_BF16" != "0" ]; then \
|
||||
echo "ERROR: Cannot use ARM-specific ISA flags (ARM_BF16) when building for x86_64 (--platform=linux/amd64)"; \
|
||||
exit 1; \
|
||||
fi
|
||||
|
||||
# Copy build requirements
|
||||
COPY requirements/cpu-build.txt requirements/build.txt
|
||||
|
||||
@@ -224,6 +238,7 @@ ARG VLLM_CPU_AVX512
|
||||
ARG VLLM_CPU_AVX512BF16
|
||||
ARG VLLM_CPU_AVX512VNNI
|
||||
ARG VLLM_CPU_AMXBF16
|
||||
ARG VLLM_CPU_ARM_BF16
|
||||
ARG PYTHON_VERSION
|
||||
|
||||
LABEL ai.vllm.build.target-arch="${TARGETARCH}"
|
||||
@@ -233,6 +248,7 @@ LABEL ai.vllm.build.cpu-avx512="${VLLM_CPU_AVX512:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx512bf16="${VLLM_CPU_AVX512BF16:-false}"
|
||||
LABEL ai.vllm.build.cpu-avx512vnni="${VLLM_CPU_AVX512VNNI:-false}"
|
||||
LABEL ai.vllm.build.cpu-amxbf16="${VLLM_CPU_AMXBF16:-false}"
|
||||
LABEL ai.vllm.build.cpu-arm-bf16="${VLLM_CPU_ARM_BF16:-false}"
|
||||
LABEL ai.vllm.build.python-version="${PYTHON_VERSION:-3.12}"
|
||||
|
||||
ENTRYPOINT ["vllm", "serve"]
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 325 KiB After Width: | Height: | Size: 331 KiB |
File diff suppressed because one or more lines are too long
|
After Width: | Height: | Size: 339 KiB |
File diff suppressed because one or more lines are too long
|
After Width: | Height: | Size: 374 KiB |
@@ -182,6 +182,7 @@ The following table lists backends that support full CUDA Graphs at the time of
|
||||
| FlashInfer | `UNIFORM_SINGLE_TOKEN_DECODE` | Will be set to `UNIFORM_BATCH` when using TRTLLM attention on Blackwell |
|
||||
| FlashMLA | `UNIFORM_BATCH` | |
|
||||
| FlashInferMLA | `UNIFORM_BATCH` | |
|
||||
| FlashInferMLASparse | `UNIFORM_BATCH` | |
|
||||
| AITER MLA | `UNIFORM_SINGLE_TOKEN_DECODE` | |
|
||||
| CUTLASS MLA | `UNIFORM_SINGLE_TOKEN_DECODE` | |
|
||||
| Mamba attention| `UNIFORM_SINGLE_TOKEN_DECODE` | |
|
||||
|
||||
@@ -155,3 +155,4 @@ The interface for the model/module may change during vLLM's development. If you
|
||||
- `use_v1` parameter in `Platform.get_attn_backend_cls` is deprecated. It has been removed in v0.13.0.
|
||||
- `_Backend` in `vllm.attention` is deprecated. It has been removed in v0.13.0. Please use `vllm.v1.attention.backends.registry.register_backend` to add new attention backend to `AttentionBackendEnum` instead.
|
||||
- `seed_everything` platform interface is deprecated. It has been removed in v0.16.0. Please use `vllm.utils.torch_utils.set_random_seed` instead.
|
||||
- `prompt` in `Platform.validate_request` is deprecated and will be removed in v0.18.0.
|
||||
|
||||
@@ -36,12 +36,12 @@ th:not(:first-child) {
|
||||
}
|
||||
</style>
|
||||
|
||||
| Feature | [CP](../configuration/optimization.md#chunked-prefill) | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](spec_decode/README.md) | CUDA graph | [pooling](../models/pooling_models.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search | [prompt-embeds](prompt_embeds.md) |
|
||||
| Feature | [CP](../configuration/optimization.md#chunked-prefill) | [APC](automatic_prefix_caching.md) | [LoRA](lora.md) | [SD](speculative_decoding/README.md) | CUDA graph | [pooling](../models/pooling_models.md) | <abbr title="Encoder-Decoder Models">enc-dec</abbr> | <abbr title="Logprobs">logP</abbr> | <abbr title="Prompt Logprobs">prmpt logP</abbr> | <abbr title="Async Output Processing">async output</abbr> | multi-step | <abbr title="Multimodal Inputs">mm</abbr> | best-of | beam-search | [prompt-embeds](prompt_embeds.md) |
|
||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
||||
| [CP](../configuration/optimization.md#chunked-prefill) | ✅ | | | | | | | | | | | | | | |
|
||||
| [APC](automatic_prefix_caching.md) | ✅ | ✅ | | | | | | | | | | | | | |
|
||||
| [LoRA](lora.md) | ✅ | ✅ | ✅ | | | | | | | | | | | | |
|
||||
| [SD](spec_decode/README.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | | |
|
||||
| [SD](speculative_decoding/README.md) | ✅ | ✅ | ❌ | ✅ | | | | | | | | | | | |
|
||||
| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
|
||||
| [pooling](../models/pooling_models.md) | 🟠\* | 🟠\* | ✅ | ❌ | ✅ | ✅ | | | | | | | | | |
|
||||
| <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ❌ | [❌](https://github.com/vllm-project/vllm/issues/7366) | ❌ | [❌](https://github.com/vllm-project/vllm/issues/7366) | ✅ | ✅ | ✅ | | | | | | | | |
|
||||
@@ -64,7 +64,7 @@ th:not(:first-child) {
|
||||
| [CP](../configuration/optimization.md#chunked-prefill) | [❌](https://github.com/vllm-project/vllm/issues/2729) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| [APC](automatic_prefix_caching.md) | [❌](https://github.com/vllm-project/vllm/issues/3687) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| [LoRA](lora.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| [SD](spec_decode/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ |
|
||||
| [SD](speculative_decoding/README.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ |
|
||||
| CUDA graph | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | [❌](https://github.com/vllm-project/vllm/issues/26970) |
|
||||
| [pooling](../models/pooling_models.md) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| <abbr title="Encoder-Decoder Models">enc-dec</abbr> | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ |
|
||||
|
||||
@@ -109,6 +109,7 @@ Batch invariance has been tested and verified on the following models:
|
||||
- **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
|
||||
- **Llama 3**: `meta-llama/Llama-3.1-8B-Instruct`, `meta-llama/Llama-3.2-1B-Instruct`
|
||||
- **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b`
|
||||
- **Mistral**: `mistralai/Mistral-7B-v0.3`
|
||||
|
||||
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).
|
||||
|
||||
|
||||
@@ -84,7 +84,7 @@ Since simple RTN does not require data for weight quantization and the activatio
|
||||
Install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.9.2"
|
||||
pip install vllm "lm-eval[api]>=0.4.11"
|
||||
```
|
||||
|
||||
Load and run the model in `vllm`:
|
||||
|
||||
@@ -18,7 +18,7 @@ pip install llmcompressor
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.9.2"
|
||||
pip install vllm "lm-eval[api]>=0.4.11"
|
||||
```
|
||||
|
||||
## Quantization Process
|
||||
|
||||
@@ -23,7 +23,7 @@ pip install llmcompressor
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.9.2"
|
||||
pip install vllm "lm-eval[api]>=0.4.11"
|
||||
```
|
||||
|
||||
## Quantization Process
|
||||
|
||||
@@ -20,7 +20,7 @@ for more installation details.
|
||||
Additionally, install `vllm` and `lm-evaluation-harness` for evaluation:
|
||||
|
||||
```bash
|
||||
pip install vllm "lm-eval[api]>=0.4.9.2"
|
||||
pip install vllm "lm-eval[api]>=0.4.11"
|
||||
```
|
||||
|
||||
## Quantization Process
|
||||
|
||||
@@ -1,330 +0,0 @@
|
||||
# Speculative Decoding
|
||||
|
||||
!!! warning
|
||||
Currently, speculative decoding in vLLM is not compatible with pipeline parallelism.
|
||||
|
||||
This document shows how to use [Speculative Decoding](https://x.com/karpathy/status/1697318534555336961) with vLLM.
|
||||
Speculative decoding is a technique which improves inter-token latency in memory-bound LLM inference.
|
||||
|
||||
!!! tip
|
||||
To train your own draft models for speculative decoding, see [Speculators](speculators.md), a library for training draft models that integrates seamlessly with vLLM.
|
||||
|
||||
## Speculating with a draft model
|
||||
|
||||
The following code configures vLLM in an offline mode to use speculative decoding with a draft model, speculating 5 tokens at a time.
|
||||
|
||||
!!! warning
|
||||
In vllm v0.10.0, speculative decoding with a draft model is not supported.
|
||||
If you use the following code, you will get a `NotImplementedError`.
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = [
|
||||
"The future of AI is",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="facebook/opt-6.7b",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"model": "facebook/opt-125m",
|
||||
"num_speculative_tokens": 5,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
To perform the same with an online mode launch the server:
|
||||
|
||||
```bash
|
||||
vllm serve facebook/opt-6.7b \
|
||||
--host 0.0.0.0 \
|
||||
--port 8000 \
|
||||
--seed 42 \
|
||||
-tp 1 \
|
||||
--gpu_memory_utilization 0.8 \
|
||||
--speculative_config '{"model": "facebook/opt-125m", "num_speculative_tokens": 5}'
|
||||
```
|
||||
|
||||
!!! warning
|
||||
Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated now.
|
||||
|
||||
Then use a client:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
# Modify OpenAI's API key and API base to use vLLM's API server.
|
||||
openai_api_key = "EMPTY"
|
||||
openai_api_base = "http://localhost:8000/v1"
|
||||
|
||||
client = OpenAI(
|
||||
# defaults to os.environ.get("OPENAI_API_KEY")
|
||||
api_key=openai_api_key,
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
|
||||
models = client.models.list()
|
||||
model = models.data[0].id
|
||||
|
||||
# Completion API
|
||||
stream = False
|
||||
completion = client.completions.create(
|
||||
model=model,
|
||||
prompt="The future of AI is",
|
||||
echo=False,
|
||||
n=1,
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
print("Completion results:")
|
||||
if stream:
|
||||
for c in completion:
|
||||
print(c)
|
||||
else:
|
||||
print(completion)
|
||||
```
|
||||
|
||||
## Speculating by matching n-grams in the prompt
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by
|
||||
matching n-grams in the prompt. For more information read [this thread.](https://x.com/joao_gante/status/1747322413006643259)
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = [
|
||||
"The future of AI is",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="facebook/opt-6.7b",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"method": "ngram",
|
||||
"num_speculative_tokens": 5,
|
||||
"prompt_lookup_max": 4,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
## Speculating using Suffix Decoding
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated using Suffix Decoding ([technical report](https://arxiv.org/abs/2411.04975)).
|
||||
|
||||
Like n-gram, Suffix Decoding can generate draft tokens by pattern-matching using the last `n` generated tokens. Unlike n-gram, Suffix Decoding (1) can pattern-match against both the prompt and previous generations, (2) uses frequency counts to propose the most likely continuations, and (3) speculates an adaptive number of tokens for each request at each iteration to get better acceptance rates.
|
||||
|
||||
Suffix Decoding can achieve better performance for tasks with high repetition, such as code-editing, agentic loops (e.g. self-reflection, self-consistency), and RL rollouts.
|
||||
|
||||
!!! tip "Install Arctic Inference"
|
||||
Suffix Decoding requires [Arctic Inference](https://github.com/snowflakedb/ArcticInference). You can install it with `pip install arctic-inference`.
|
||||
|
||||
!!! tip "Suffix Decoding Speculative Tokens"
|
||||
Suffix Decoding will speculate a dynamic number of tokens for each request at each decoding step, so the `num_speculative_tokens` configuration specifies the *maximum* number of speculative tokens. It is suggested to use a high number such as `16` or `32` (default).
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = [
|
||||
"The future of AI is",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="facebook/opt-6.7b",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"method": "suffix",
|
||||
"num_speculative_tokens": 32,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
## Speculating using MLP speculators
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by
|
||||
draft models that condition draft predictions on both context vectors and sampled tokens.
|
||||
For more information see [this blog](https://pytorch.org/blog/hitchhikers-guide-speculative-decoding/) or
|
||||
[this technical report](https://arxiv.org/abs/2404.19124).
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = [
|
||||
"The future of AI is",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3.1-70B-Instruct",
|
||||
tensor_parallel_size=4,
|
||||
speculative_config={
|
||||
"model": "ibm-ai-platform/llama3-70b-accelerator",
|
||||
"draft_tensor_parallel_size": 1,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
Note that these speculative models currently need to be run without tensor parallelism, although
|
||||
it is possible to run the main model using tensor parallelism (see example above). Since the
|
||||
speculative models are relatively small, we still see significant speedups. However, this
|
||||
limitation will be fixed in a future release.
|
||||
|
||||
A variety of speculative models of this type are available on HF hub:
|
||||
|
||||
- [llama-13b-accelerator](https://huggingface.co/ibm-ai-platform/llama-13b-accelerator)
|
||||
- [llama3-8b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-8b-accelerator)
|
||||
- [codellama-34b-accelerator](https://huggingface.co/ibm-ai-platform/codellama-34b-accelerator)
|
||||
- [llama2-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama2-70b-accelerator)
|
||||
- [llama3-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-70b-accelerator)
|
||||
- [granite-3b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-3b-code-instruct-accelerator)
|
||||
- [granite-8b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-8b-code-instruct-accelerator)
|
||||
- [granite-7b-instruct-accelerator](https://huggingface.co/ibm-granite/granite-7b-instruct-accelerator)
|
||||
- [granite-20b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-20b-code-instruct-accelerator)
|
||||
|
||||
## Speculating using EAGLE based draft models
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by
|
||||
an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/offline_inference/spec_decode.py](../../../examples/offline_inference/spec_decode.py)
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = [
|
||||
"The future of AI is",
|
||||
]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3-8B-Instruct",
|
||||
tensor_parallel_size=4,
|
||||
speculative_config={
|
||||
"model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
|
||||
"draft_tensor_parallel_size": 1,
|
||||
"num_speculative_tokens": 2,
|
||||
"method": "eagle",
|
||||
},
|
||||
)
|
||||
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
|
||||
```
|
||||
|
||||
A few important things to consider when using the EAGLE based draft models:
|
||||
|
||||
1. The EAGLE draft models available in the [HF repository for EAGLE models](https://huggingface.co/yuhuili) should
|
||||
be able to be loaded and used directly by vLLM after <https://github.com/vllm-project/vllm/pull/12304>.
|
||||
If you are using vllm version before <https://github.com/vllm-project/vllm/pull/12304>, please use the
|
||||
[script](https://gist.github.com/abhigoyal1997/1e7a4109ccb7704fbc67f625e86b2d6d) to convert the speculative model,
|
||||
and specify `"model": "path/to/modified/eagle/model"` in `speculative_config`. If weight-loading problems still occur when using the latest version of vLLM, please leave a comment or raise an issue.
|
||||
|
||||
2. The EAGLE based draft models need to be run without tensor parallelism
|
||||
(i.e. draft_tensor_parallel_size is set to 1 in `speculative_config`), although
|
||||
it is possible to run the main model using tensor parallelism (see example above).
|
||||
|
||||
3. When using EAGLE-based speculators with vLLM, the observed speedup is lower than what is
|
||||
reported in the reference implementation [here](https://github.com/SafeAILab/EAGLE). This issue is under
|
||||
investigation and tracked here: <https://github.com/vllm-project/vllm/issues/9565>.
|
||||
|
||||
4. When using EAGLE-3 based draft model, option "method" must be set to "eagle3".
|
||||
That is, to specify `"method": "eagle3"` in `speculative_config`.
|
||||
|
||||
A variety of EAGLE draft models are available on the Hugging Face hub:
|
||||
|
||||
| Base Model | EAGLE on Hugging Face | # EAGLE Parameters |
|
||||
|---------------------------------------------------------------------|-------------------------------------------|--------------------|
|
||||
| Vicuna-7B-v1.3 | yuhuili/EAGLE-Vicuna-7B-v1.3 | 0.24B |
|
||||
| Vicuna-13B-v1.3 | yuhuili/EAGLE-Vicuna-13B-v1.3 | 0.37B |
|
||||
| Vicuna-33B-v1.3 | yuhuili/EAGLE-Vicuna-33B-v1.3 | 0.56B |
|
||||
| LLaMA2-Chat 7B | yuhuili/EAGLE-llama2-chat-7B | 0.24B |
|
||||
| LLaMA2-Chat 13B | yuhuili/EAGLE-llama2-chat-13B | 0.37B |
|
||||
| LLaMA2-Chat 70B | yuhuili/EAGLE-llama2-chat-70B | 0.99B |
|
||||
| Mixtral-8x7B-Instruct-v0.1 | yuhuili/EAGLE-mixtral-instruct-8x7B | 0.28B |
|
||||
| LLaMA3-Instruct 8B | yuhuili/EAGLE-LLaMA3-Instruct-8B | 0.25B |
|
||||
| LLaMA3-Instruct 70B | yuhuili/EAGLE-LLaMA3-Instruct-70B | 0.99B |
|
||||
| Qwen2-7B-Instruct | yuhuili/EAGLE-Qwen2-7B-Instruct | 0.26B |
|
||||
| Qwen2-72B-Instruct | yuhuili/EAGLE-Qwen2-72B-Instruct | 1.05B |
|
||||
|
||||
## Lossless guarantees of Speculative Decoding
|
||||
|
||||
In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
|
||||
speculative decoding, breaking down the guarantees into three key areas:
|
||||
|
||||
1. **Theoretical Losslessness**
|
||||
\- Speculative decoding sampling is theoretically lossless up to the precision limits of hardware numerics. Floating-point errors might
|
||||
cause slight variations in output distributions, as discussed
|
||||
in [Accelerating Large Language Model Decoding with Speculative Sampling](https://arxiv.org/pdf/2302.01318)
|
||||
|
||||
2. **Algorithmic Losslessness**
|
||||
\- vLLM’s implementation of speculative decoding is algorithmically validated to be lossless. Key validation tests include:
|
||||
|
||||
> - **Rejection Sampler Convergence**: Ensures that samples from vLLM’s rejection sampler align with the target
|
||||
> distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
|
||||
> - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
|
||||
> without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
|
||||
> provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](../../tests/spec_decode/e2e).
|
||||
> verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
|
||||
|
||||
3. **vLLM Logprob Stability**
|
||||
\- vLLM does not currently guarantee stable token log probabilities (logprobs). This can result in different outputs for the
|
||||
same request across runs. For more details, see the FAQ section
|
||||
titled *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
|
||||
|
||||
While vLLM strives to ensure losslessness in speculative decoding, variations in generated outputs with and without speculative decoding
|
||||
can occur due to following factors:
|
||||
|
||||
- **Floating-Point Precision**: Differences in hardware numerical precision may lead to slight discrepancies in the output distribution.
|
||||
- **Batch Size and Numerical Stability**: Changes in batch size may cause variations in logprobs and output probabilities, potentially
|
||||
due to non-deterministic behavior in batched operations or numerical instability.
|
||||
|
||||
For mitigation strategies, please refer to the FAQ entry *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
|
||||
|
||||
## Resources for vLLM contributors
|
||||
|
||||
- [A Hacker's Guide to Speculative Decoding in vLLM](https://www.youtube.com/watch?v=9wNAgpX6z_4)
|
||||
- [What is Lookahead Scheduling in vLLM?](https://docs.google.com/document/d/1Z9TvqzzBPnh5WHcRwjvK2UEeFeq5zMZb5mFE8jR0HCs/edit#heading=h.1fjfb0donq5a)
|
||||
- [Information on batch expansion](https://docs.google.com/document/d/1T-JaS2T1NRfdP51qzqpyakoCXxSXTtORppiwaj5asxA/edit#heading=h.kk7dq05lc6q8)
|
||||
- [Dynamic speculative decoding](https://github.com/vllm-project/vllm/issues/4565)
|
||||
@@ -0,0 +1,62 @@
|
||||
# Speculative Decoding
|
||||
|
||||
This document shows how to use [Speculative Decoding](https://arxiv.org/pdf/2302.01318) with vLLM to reduce inter-token latency under medium-to-low QPS (query per second), memory-bound workloads.
|
||||
|
||||
To train your own draft models for optimized speculative decoding, see [vllm-project/speculators](speculators.md) for seamless training and integration with vLLM.
|
||||
|
||||
## vLLM Speculation Methods
|
||||
|
||||
vLLM supports a variety of methods of speculative decoding. Model-based methods such as EAGLE, draft models, and mlp provide the best latency reduction, while simpler methods such as n-gram and and suffix decoding provide modest speedups without increasing workload during peak traffic.
|
||||
|
||||
- [EAGLE](eagle.md)
|
||||
- [Draft Model](draft_model.md)
|
||||
- [Multi-Layer Perceptron](mlp.md)
|
||||
- [N-Gram](n_gram.md)
|
||||
- [Suffix Decoding](suffix.md)
|
||||
|
||||
## Lossless guarantees of Speculative Decoding
|
||||
|
||||
In vLLM, speculative decoding aims to enhance inference efficiency while maintaining accuracy. This section addresses the lossless guarantees of
|
||||
speculative decoding, breaking down the guarantees into three key areas:
|
||||
|
||||
1. **Theoretical Losslessness**
|
||||
\- Speculative decoding sampling is theoretically lossless up to the precision limits of hardware numerics. Floating-point errors might
|
||||
cause slight variations in output distributions, as discussed
|
||||
in [Accelerating Large Language Model Decoding with Speculative Sampling](https://arxiv.org/pdf/2302.01318)
|
||||
|
||||
2. **Algorithmic Losslessness**
|
||||
\- vLLM’s implementation of speculative decoding is algorithmically validated to be lossless. Key validation tests include:
|
||||
|
||||
> - **Rejection Sampler Convergence**: Ensures that samples from vLLM’s rejection sampler align with the target
|
||||
> distribution. [View Test Code](https://github.com/vllm-project/vllm/blob/47b65a550866c7ffbd076ecb74106714838ce7da/tests/samplers/test_rejection_sampler.py#L252)
|
||||
> - **Greedy Sampling Equality**: Confirms that greedy sampling with speculative decoding matches greedy sampling
|
||||
> without it. This verifies that vLLM's speculative decoding framework, when integrated with the vLLM forward pass and the vLLM rejection sampler,
|
||||
> provides a lossless guarantee. Almost all of the tests in [tests/spec_decode/e2e](/tests/v1/spec_decode).
|
||||
> verify this property using [this assertion implementation](https://github.com/vllm-project/vllm/blob/b67ae00cdbbe1a58ffc8ff170f0c8d79044a684a/tests/spec_decode/e2e/conftest.py#L291)
|
||||
|
||||
3. **vLLM Logprob Stability**
|
||||
\- vLLM does not currently guarantee stable token log probabilities (logprobs). This can result in different outputs for the
|
||||
same request across runs. For more details, see the FAQ section
|
||||
titled *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
|
||||
|
||||
While vLLM strives to ensure losslessness in speculative decoding, variations in generated outputs with and without speculative decoding
|
||||
can occur due to following factors:
|
||||
|
||||
- **Floating-Point Precision**: Differences in hardware numerical precision may lead to slight discrepancies in the output distribution.
|
||||
- **Batch Size and Numerical Stability**: Changes in batch size may cause variations in logprobs and output probabilities, potentially
|
||||
due to non-deterministic behavior in batched operations or numerical instability.
|
||||
|
||||
For mitigation strategies, please refer to the FAQ entry *Can the output of a prompt vary across runs in vLLM?* in the [FAQs](../../usage/faq.md).
|
||||
|
||||
## Known Feature Incompatibility
|
||||
|
||||
1. Pipeline parallelism is not composible with speculative decoding as of `vllm<=0.15.0`
|
||||
2. Speculative decoding with a draft models is not supported in `vllm<=0.10.0`
|
||||
|
||||
## Resources for vLLM contributors
|
||||
|
||||
- [[vLLM Office Hours #40] Intro to Speculators](https://www.youtube.com/watch?v=2ISAr_JVGLs)
|
||||
- [A Hacker's Guide to Speculative Decoding in vLLM](https://www.youtube.com/watch?v=9wNAgpX6z_4)
|
||||
- [What is Lookahead Scheduling in vLLM?](https://docs.google.com/document/d/1Z9TvqzzBPnh5WHcRwjvK2UEeFeq5zMZb5mFE8jR0HCs/edit#heading=h.1fjfb0donq5a)
|
||||
- [Information on batch expansion](https://docs.google.com/document/d/1T-JaS2T1NRfdP51qzqpyakoCXxSXTtORppiwaj5asxA/edit#heading=h.kk7dq05lc6q8)
|
||||
- [Dynamic speculative decoding](https://github.com/vllm-project/vllm/issues/4565)
|
||||
@@ -0,0 +1,80 @@
|
||||
# Draft Models
|
||||
|
||||
The following code configures vLLM in an offline mode to use speculative decoding with a draft model, speculating 5 tokens at a time.
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-8B",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"model": "Qwen/Qwen3-0.6B",
|
||||
"num_speculative_tokens": 5,
|
||||
"method": "draft_model",
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
To perform the equivalent launch in online mode, use the following server-side code:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-4B-Thinking-2507 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8000 \
|
||||
--seed 42 \
|
||||
-tp 1 \
|
||||
--max_model_len 2048 \
|
||||
--gpu_memory_utilization 0.8 \
|
||||
--speculative_config '{"model": "Qwen/Qwen3-0.6B", "num_speculative_tokens": 5, "method": "draft_model"}'
|
||||
```
|
||||
|
||||
The code used to request as completions as a client remains unchanged:
|
||||
|
||||
??? code
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
# Modify OpenAI's API key and API base to use vLLM's API server.
|
||||
openai_api_key = "EMPTY"
|
||||
openai_api_base = "http://localhost:8000/v1"
|
||||
|
||||
client = OpenAI(
|
||||
# defaults to os.environ.get("OPENAI_API_KEY")
|
||||
api_key=openai_api_key,
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
|
||||
models = client.models.list()
|
||||
model = models.data[0].id
|
||||
|
||||
# Completion API
|
||||
stream = False
|
||||
completion = client.completions.create(
|
||||
model=model,
|
||||
prompt="The future of AI is",
|
||||
echo=False,
|
||||
n=1,
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
print("Completion results:")
|
||||
if stream:
|
||||
for c in completion:
|
||||
print(c)
|
||||
else:
|
||||
print(completion)
|
||||
```
|
||||
|
||||
!!! warning
|
||||
Note: Please use `--speculative_config` to set all configurations related to speculative decoding. The previous method of specifying the model through `--speculative_model` and adding related parameters (e.g., `--num_speculative_tokens`) separately has been deprecated.
|
||||
@@ -0,0 +1,67 @@
|
||||
# EAGLE Draft Models
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by an [EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency)](https://arxiv.org/pdf/2401.15077) based draft model. A more detailed example for offline mode, including how to extract request level acceptance rate, can be found in [examples/offline_inference/spec_decode.py](../../../examples/offline_inference/spec_decode.py)
|
||||
|
||||
## Eagle Drafter Example
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3-8B-Instruct",
|
||||
tensor_parallel_size=4,
|
||||
speculative_config={
|
||||
"model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
|
||||
"draft_tensor_parallel_size": 1,
|
||||
"num_speculative_tokens": 2,
|
||||
"method": "eagle",
|
||||
},
|
||||
)
|
||||
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
## Eagle3 Drafter Example
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3-8B-Instruct",
|
||||
tensor_parallel_size=2,
|
||||
speculative_config={
|
||||
"model": "RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3",
|
||||
"draft_tensor_parallel_size": 2,
|
||||
"num_speculative_tokens": 2,
|
||||
"method": "eagle",
|
||||
},
|
||||
)
|
||||
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
## Pre-Trained Eagle Draft Models
|
||||
|
||||
A variety of EAGLE draft models are available on the Hugging Face hub:
|
||||
|
||||
* [RedHatAI/speculator-models](https://huggingface.co/collections/RedHatAI/speculator-models)
|
||||
* [yuhuili/models](https://huggingface.co/yuhuili/models?search=eagle)
|
||||
|
||||
!!! warning
|
||||
If you are using `vllm<0.7.0`, please use [this script](https://gist.github.com/abhigoyal1997/1e7a4109ccb7704fbc67f625e86b2d6d) to convert the speculative model and specify `"model": "path/to/modified/eagle/model"` in `speculative_config`.
|
||||
@@ -0,0 +1,42 @@
|
||||
# MLP Draft Models
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by draft models that condition draft predictions on both context vectors and sampled tokens. For more information see [The Hitchhiker's Guide to Speculative Decoding](https://pytorch.org/blog/hitchhikers-guide-speculative-decoding/) and [IBM Research's Technical Report](https://arxiv.org/abs/2404.19124).
|
||||
|
||||
## MLP Drafter Example
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="meta-llama/Meta-Llama-3.1-70B-Instruct",
|
||||
tensor_parallel_size=4,
|
||||
speculative_config={
|
||||
"model": "ibm-ai-platform/llama3-70b-accelerator",
|
||||
"draft_tensor_parallel_size": 1,
|
||||
"method": "mlp_speculator",
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
|
||||
## Pre-Trained MLP Drafter Models
|
||||
|
||||
A variety of speculative models of this type are available on HF hub:
|
||||
|
||||
- [llama-13b-accelerator](https://huggingface.co/ibm-ai-platform/llama-13b-accelerator)
|
||||
- [llama3-8b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-8b-accelerator)
|
||||
- [codellama-34b-accelerator](https://huggingface.co/ibm-ai-platform/codellama-34b-accelerator)
|
||||
- [llama2-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama2-70b-accelerator)
|
||||
- [llama3-70b-accelerator](https://huggingface.co/ibm-ai-platform/llama3-70b-accelerator)
|
||||
- [granite-3b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-3b-code-instruct-accelerator)
|
||||
- [granite-8b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-8b-code-instruct-accelerator)
|
||||
- [granite-7b-instruct-accelerator](https://huggingface.co/ibm-granite/granite-7b-instruct-accelerator)
|
||||
- [granite-20b-code-instruct-accelerator](https://huggingface.co/ibm-granite/granite-20b-code-instruct-accelerator)
|
||||
@@ -0,0 +1,27 @@
|
||||
# N-Gram Speculation
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated by
|
||||
matching n-grams in the prompt. For more information read [this thread.](https://x.com/joao_gante/status/1747322413006643259)
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-8B",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"method": "ngram",
|
||||
"num_speculative_tokens": 5,
|
||||
"prompt_lookup_max": 4,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
+4
-1
@@ -1,4 +1,7 @@
|
||||
# Speculators
|
||||
# vLLM-Project/Speculators
|
||||
|
||||

|
||||

|
||||
|
||||
[Speculators](https://docs.vllm.ai/projects/speculators/en/latest/) is a library for accelerating LLM inference through speculative decoding, providing efficient draft model training that integrates seamlessly with vLLM to reduce latency and improve throughput.
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
# Suffix Decoding
|
||||
|
||||
The following code configures vLLM to use speculative decoding where proposals are generated using Suffix Decoding ([technical report](https://arxiv.org/abs/2411.04975)).
|
||||
|
||||
Like n-gram, Suffix Decoding can generate draft tokens by pattern-matching using the last `n` generated tokens. Unlike n-gram, Suffix Decoding (1) can pattern-match against both the prompt and previous generations, (2) uses frequency counts to propose the most likely continuations, and (3) speculates an adaptive number of tokens for each request at each iteration to get better acceptance rates.
|
||||
|
||||
Suffix Decoding can achieve better performance for tasks with high repetition, such as code-editing, agentic loops (e.g. self-reflection, self-consistency), and RL rollouts.
|
||||
|
||||
!!! tip "Install Arctic Inference"
|
||||
Suffix Decoding requires [Arctic Inference](https://github.com/snowflakedb/ArcticInference). You can install it with `pip install arctic-inference`.
|
||||
|
||||
!!! tip "Suffix Decoding Speculative Tokens"
|
||||
Suffix Decoding will speculate a dynamic number of tokens for each request at each decoding step, so the `num_speculative_tokens` configuration specifies the *maximum* number of speculative tokens. It is suggested to use a high number such as `16` or `32` (default).
|
||||
|
||||
```python
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
prompts = ["The future of AI is"]
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-8B",
|
||||
tensor_parallel_size=1,
|
||||
speculative_config={
|
||||
"method": "suffix",
|
||||
"num_speculative_tokens": 32,
|
||||
},
|
||||
)
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
```
|
||||
@@ -172,25 +172,78 @@ docker pull public.ecr.aws/q9t5s3a7/vllm-ci-postmerge-repo:${VLLM_COMMIT}-arm64-
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
# --8<-- [start:build-image-from-source]
|
||||
|
||||
## Building for your target ARM CPU
|
||||
|
||||
```bash
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--tag vllm-cpu-env .
|
||||
--platform=linux/arm64 \
|
||||
--build-arg VLLM_CPU_ARM_BF16=<false (default)|true> \
|
||||
--tag vllm-cpu-env \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
# Launching OpenAI server
|
||||
!!! note "Auto-detection by default"
|
||||
By default, ARM CPU instruction sets (BF16, NEON, etc.) are automatically detected from the build system's CPU flags. The `VLLM_CPU_ARM_BF16` build argument is used for cross-compilation:
|
||||
|
||||
- `VLLM_CPU_ARM_BF16=true` - Force-enable ARM BF16 support (build with BF16 regardless of build system capabilities)
|
||||
- `VLLM_CPU_ARM_BF16=false` - Rely on auto-detection (default)
|
||||
|
||||
### Examples
|
||||
|
||||
**Auto-detection build (native ARM)**
|
||||
|
||||
```bash
|
||||
# Building on ARM64 system - platform auto-detected
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--tag vllm-cpu-arm64 \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
**Cross-compile for ARM with BF16 support**
|
||||
|
||||
```bash
|
||||
# Building on ARM64 for newer ARM CPUs with BF16
|
||||
docker build -f docker/Dockerfile.cpu \
|
||||
--build-arg VLLM_CPU_ARM_BF16=true \
|
||||
--tag vllm-cpu-arm64-bf16 \
|
||||
--target vllm-openai .
|
||||
```
|
||||
|
||||
**Cross-compile from x86_64 to ARM64 with BF16**
|
||||
|
||||
```bash
|
||||
# Requires Docker buildx with ARM emulation (QEMU)
|
||||
docker buildx build -f docker/Dockerfile.cpu \
|
||||
--platform=linux/arm64 \
|
||||
--build-arg VLLM_CPU_ARM_BF16=true \
|
||||
--build-arg max_jobs=4 \
|
||||
--tag vllm-cpu-arm64-bf16 \
|
||||
--target vllm-openai \
|
||||
--load .
|
||||
```
|
||||
|
||||
!!! note "ARM BF16 requirements"
|
||||
ARM BF16 support requires ARMv8.6-A or later (FEAT_BF16). Supported on AWS Graviton3/4, AmpereOne, and other recent ARM processors.
|
||||
|
||||
## Launching the OpenAI server
|
||||
|
||||
```bash
|
||||
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 \
|
||||
vllm-cpu-arm64 \
|
||||
meta-llama/Llama-3.2-1B-Instruct \
|
||||
--dtype=bfloat16 \
|
||||
other vLLM OpenAI server arguments
|
||||
```
|
||||
|
||||
!!! tip
|
||||
An alternative of `--privileged=true` is `--cap-add SYS_NICE --security-opt seccomp=unconfined`.
|
||||
!!! tip "Alternative to --privileged"
|
||||
Instead of `--privileged=true`, use `--cap-add SYS_NICE --security-opt seccomp=unconfined` for better security.
|
||||
|
||||
# --8<-- [end:build-image-from-source]
|
||||
# --8<-- [start:extra-information]
|
||||
|
||||
@@ -181,6 +181,7 @@ Some model architectures are supported via vLLM plugins. These plugins extend vL
|
||||
| Architecture | Models | Plugin Repository |
|
||||
|--------------|--------|-------------------|
|
||||
| `BartForConditionalGeneration` | BART | [bart-plugin](https://github.com/vllm-project/bart-plugin) |
|
||||
| `Florence2ForConditionalGeneration` | Florence-2 | [bart-plugin](https://github.com/vllm-project/bart-plugin) |
|
||||
|
||||
For other model architectures not natively supported, in particular for Encoder-Decoder models, we recommend following a similar pattern by implementing support through the plugin system.
|
||||
|
||||
|
||||
@@ -137,6 +137,7 @@ Please note that prefix caching is not yet supported for any of the above models
|
||||
Whisper is supported natively. Other encoder-decoder models are supported via the plugin system:
|
||||
|
||||
- **BART**: `BartForConditionalGeneration` is supported via the official [bart-plugin](https://github.com/vllm-project/bart-plugin).
|
||||
- **Florence-2**: `Florence2ForConditionalGeneration` is supported via the official [bart-plugin](https://github.com/vllm-project/bart-plugin).
|
||||
|
||||
For other encoder-decoder models (e.g., `MllamaForConditionalGeneration`), we recommend
|
||||
following a similar pattern by implementing support through the [plugin system](../design/plugin_system.md).
|
||||
|
||||
@@ -8,7 +8,7 @@ declare -a PIDS=()
|
||||
###############################################################################
|
||||
MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
|
||||
LOG_PATH="${LOG_PATH:-./logs}"
|
||||
mkdir -p $LOG_PATH
|
||||
mkdir -p "$LOG_PATH"
|
||||
|
||||
ENCODE_PORT="${ENCODE_PORT:-19534}"
|
||||
PREFILL_PORT="${PREFILL_PORT:-19535}"
|
||||
@@ -84,10 +84,10 @@ trap cleanup TERM
|
||||
|
||||
# clear previous cache
|
||||
echo "remove previous ec cache folder"
|
||||
rm -rf $EC_SHARED_STORAGE_PATH
|
||||
rm -rf "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
echo "make ec cache folder"
|
||||
mkdir -p $EC_SHARED_STORAGE_PATH
|
||||
mkdir -p "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
###############################################################################
|
||||
# Encoder worker
|
||||
@@ -100,7 +100,7 @@ CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
|
||||
--no-enable-prefix-caching \
|
||||
--max-num-batched-tokens 114688 \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_producer",
|
||||
@@ -124,7 +124,7 @@ vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--enable-request-id-headers \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_consumer",
|
||||
@@ -152,7 +152,7 @@ vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--enable-request-id-headers \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--kv-transfer-config '{
|
||||
"kv_connector": "NixlConnector",
|
||||
"kv_role": "kv_consumer"
|
||||
@@ -162,9 +162,9 @@ vllm serve "$MODEL" \
|
||||
PIDS+=($!)
|
||||
|
||||
# Wait for workers
|
||||
wait_for_server $ENCODE_PORT
|
||||
wait_for_server $PREFILL_PORT
|
||||
wait_for_server $DECODE_PORT
|
||||
wait_for_server "$ENCODE_PORT"
|
||||
wait_for_server "$PREFILL_PORT"
|
||||
wait_for_server "$DECODE_PORT"
|
||||
|
||||
###############################################################################
|
||||
# Proxy
|
||||
@@ -179,7 +179,7 @@ python disagg_epd_proxy.py \
|
||||
|
||||
PIDS+=($!)
|
||||
|
||||
wait_for_server $PROXY_PORT
|
||||
wait_for_server "$PROXY_PORT"
|
||||
echo "All services are up!"
|
||||
|
||||
###############################################################################
|
||||
@@ -187,14 +187,14 @@ echo "All services are up!"
|
||||
###############################################################################
|
||||
echo "Running benchmark (stream)..."
|
||||
vllm bench serve \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--backend openai-chat \
|
||||
--endpoint /v1/chat/completions \
|
||||
--dataset-name hf \
|
||||
--dataset-path lmarena-ai/VisionArena-Chat \
|
||||
--seed 0 \
|
||||
--num-prompts $NUM_PROMPTS \
|
||||
--port $PROXY_PORT
|
||||
--num-prompts "$NUM_PROMPTS" \
|
||||
--port "$PROXY_PORT"
|
||||
|
||||
PIDS+=($!)
|
||||
|
||||
@@ -202,10 +202,10 @@ PIDS+=($!)
|
||||
# Single request with local image
|
||||
###############################################################################
|
||||
echo "Running single request with local image (non-stream)..."
|
||||
curl http://127.0.0.1:${PROXY_PORT}/v1/chat/completions \
|
||||
curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "'${MODEL}'",
|
||||
"model": "'"${MODEL}"'",
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": [
|
||||
|
||||
@@ -8,7 +8,7 @@ declare -a PIDS=()
|
||||
###############################################################################
|
||||
MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
|
||||
LOG_PATH="${LOG_PATH:-./logs}"
|
||||
mkdir -p $LOG_PATH
|
||||
mkdir -p "$LOG_PATH"
|
||||
|
||||
ENCODE_PORT="${ENCODE_PORT:-19534}"
|
||||
PREFILL_DECODE_PORT="${PREFILL_DECODE_PORT:-19535}"
|
||||
@@ -78,10 +78,10 @@ trap cleanup TERM
|
||||
|
||||
# clear previous cache
|
||||
echo "remove previous ec cache folder"
|
||||
rm -rf $EC_SHARED_STORAGE_PATH
|
||||
rm -rf "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
echo "make ec cache folder"
|
||||
mkdir -p $EC_SHARED_STORAGE_PATH
|
||||
mkdir -p "$EC_SHARED_STORAGE_PATH"
|
||||
|
||||
###############################################################################
|
||||
# Encoder worker
|
||||
@@ -94,7 +94,7 @@ CUDA_VISIBLE_DEVICES="$GPU_E" vllm serve "$MODEL" \
|
||||
--no-enable-prefix-caching \
|
||||
--max-num-batched-tokens 114688 \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_producer",
|
||||
@@ -115,7 +115,7 @@ CUDA_VISIBLE_DEVICES="$GPU_PD" vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--enable-request-id-headers \
|
||||
--max-num-seqs 128 \
|
||||
--allowed-local-media-path ${GIT_ROOT}/tests/v1/ec_connector/integration \
|
||||
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
|
||||
--ec-transfer-config '{
|
||||
"ec_connector": "ECExampleConnector",
|
||||
"ec_role": "ec_consumer",
|
||||
@@ -128,8 +128,8 @@ CUDA_VISIBLE_DEVICES="$GPU_PD" vllm serve "$MODEL" \
|
||||
PIDS+=($!)
|
||||
|
||||
# Wait for workers
|
||||
wait_for_server $ENCODE_PORT
|
||||
wait_for_server $PREFILL_DECODE_PORT
|
||||
wait_for_server "$ENCODE_PORT"
|
||||
wait_for_server "$PREFILL_DECODE_PORT"
|
||||
|
||||
###############################################################################
|
||||
# Proxy
|
||||
@@ -144,7 +144,7 @@ python disagg_epd_proxy.py \
|
||||
|
||||
PIDS+=($!)
|
||||
|
||||
wait_for_server $PROXY_PORT
|
||||
wait_for_server "$PROXY_PORT"
|
||||
echo "All services are up!"
|
||||
|
||||
###############################################################################
|
||||
@@ -152,14 +152,14 @@ echo "All services are up!"
|
||||
###############################################################################
|
||||
echo "Running benchmark (stream)..."
|
||||
vllm bench serve \
|
||||
--model $MODEL \
|
||||
--model "$MODEL" \
|
||||
--backend openai-chat \
|
||||
--endpoint /v1/chat/completions \
|
||||
--dataset-name hf \
|
||||
--dataset-path lmarena-ai/VisionArena-Chat \
|
||||
--seed 0 \
|
||||
--num-prompts $NUM_PROMPTS \
|
||||
--port $PROXY_PORT
|
||||
--num-prompts "$NUM_PROMPTS" \
|
||||
--port "$PROXY_PORT"
|
||||
|
||||
PIDS+=($!)
|
||||
|
||||
@@ -167,10 +167,10 @@ PIDS+=($!)
|
||||
# Single request with local image
|
||||
###############################################################################
|
||||
echo "Running single request with local image (non-stream)..."
|
||||
curl http://127.0.0.1:${PROXY_PORT}/v1/chat/completions \
|
||||
curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "'${MODEL}'",
|
||||
"model": "'"${MODEL}"'",
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": [
|
||||
|
||||
@@ -54,7 +54,7 @@ wait_for_server() {
|
||||
# You can also adjust --kv-ip and --kv-port for distributed inference.
|
||||
|
||||
# prefilling instance, which is the KV producer
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve $MODEL_NAME \
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve "$MODEL_NAME" \
|
||||
--host 0.0.0.0 \
|
||||
--port 8100 \
|
||||
--max-model-len 100 \
|
||||
@@ -64,7 +64,7 @@ CUDA_VISIBLE_DEVICES=0 vllm serve $MODEL_NAME \
|
||||
'{"kv_connector":"P2pNcclConnector","kv_role":"kv_producer","kv_rank":0,"kv_parallel_size":2,"kv_buffer_size":"1e9","kv_port":"14579","kv_connector_extra_config":{"proxy_ip":"'"$VLLM_HOST_IP"'","proxy_port":"30001","http_ip":"'"$VLLM_HOST_IP"'","http_port":"8100","send_type":"PUT_ASYNC"}}' &
|
||||
|
||||
# decoding instance, which is the KV consumer
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve $MODEL_NAME \
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve "$MODEL_NAME" \
|
||||
--host 0.0.0.0 \
|
||||
--port 8200 \
|
||||
--max-model-len 100 \
|
||||
|
||||
@@ -328,9 +328,9 @@ class Proxy:
|
||||
if instance_type == "decode" and instance in self.decode_instances:
|
||||
self.decode_instances.remove(instance)
|
||||
self.decode_cycler = itertools.cycle(self.decode_instances)
|
||||
if instance_type == "prefill" and instance in self.decode_instances:
|
||||
if instance_type == "prefill" and instance in self.prefill_instances:
|
||||
self.prefill_instances.remove(instance)
|
||||
self.prefill_cycler = itertools.cycle(self.decode_instances)
|
||||
self.prefill_cycler = itertools.cycle(self.prefill_instances)
|
||||
|
||||
|
||||
class RoundRobinSchedulingPolicy(SchedulingPolicy):
|
||||
|
||||
@@ -34,7 +34,7 @@ wait_for_server() {
|
||||
done" && return 0 || return 1
|
||||
}
|
||||
|
||||
vllm serve $MODEL_NAME \
|
||||
vllm serve "$MODEL_NAME" \
|
||||
--port 8100 \
|
||||
--max-model-len 100 \
|
||||
--enforce-eager \
|
||||
|
||||
+12
-11
@@ -143,7 +143,7 @@ main() {
|
||||
IFS=',' read -ra BOOTSTRAP_PORT_ARRAY <<< "$BOOTSTRAP_PORTS"
|
||||
IFS=',' read -ra DECODE_PORT_ARRAY <<< "$DECODE_PORTS"
|
||||
|
||||
proxy_param=""
|
||||
proxy_args=()
|
||||
|
||||
# =============================================================================
|
||||
# Launch Prefill Servers (X Producers)
|
||||
@@ -156,12 +156,12 @@ main() {
|
||||
local bootstrap_port=${BOOTSTRAP_PORT_ARRAY[$i]}
|
||||
|
||||
echo " Prefill server $((i+1)): GPU $gpu_id, Port $port, Bootstrap Port $bootstrap_port"
|
||||
VLLM_MOONCAKE_BOOTSTRAP_PORT=$bootstrap_port CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \
|
||||
--port $port \
|
||||
VLLM_MOONCAKE_BOOTSTRAP_PORT=$bootstrap_port CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
|
||||
--port "$port" \
|
||||
--kv-transfer-config \
|
||||
"{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_producer\"}" > prefill$((i+1)).log 2>&1 &
|
||||
PIDS+=($!)
|
||||
proxy_param="${proxy_param} --prefill http://0.0.0.0:${port} $bootstrap_port"
|
||||
proxy_args+=(--prefill "http://0.0.0.0:${port}" "$bootstrap_port")
|
||||
done
|
||||
|
||||
# =============================================================================
|
||||
@@ -174,12 +174,12 @@ main() {
|
||||
local port=${DECODE_PORT_ARRAY[$i]}
|
||||
|
||||
echo " Decode server $((i+1)): GPU $gpu_id, Port $port"
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \
|
||||
--port $port \
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
|
||||
--port "$port" \
|
||||
--kv-transfer-config \
|
||||
"{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"kv_consumer\"}" > decode$((i+1)).log 2>&1 &
|
||||
PIDS+=($!)
|
||||
proxy_param="${proxy_param} --decode http://0.0.0.0:${port}"
|
||||
proxy_args+=(--decode "http://0.0.0.0:${port}")
|
||||
done
|
||||
|
||||
# =============================================================================
|
||||
@@ -187,7 +187,7 @@ main() {
|
||||
# =============================================================================
|
||||
echo ""
|
||||
echo "Starting proxy server on port $PROXY_PORT..."
|
||||
python3 mooncake_connector_proxy.py $proxy_param --port $PROXY_PORT > proxy.log 2>&1 &
|
||||
python3 mooncake_connector_proxy.py "${proxy_args[@]}" --port "$PROXY_PORT" > proxy.log 2>&1 &
|
||||
PIDS+=($!)
|
||||
|
||||
# =============================================================================
|
||||
@@ -196,9 +196,10 @@ main() {
|
||||
echo ""
|
||||
echo "Waiting for all servers to start..."
|
||||
for port in "${PREFILL_PORT_ARRAY[@]}" "${DECODE_PORT_ARRAY[@]}"; do
|
||||
if ! wait_for_server $port; then
|
||||
if ! wait_for_server "$port"; then
|
||||
echo "Failed to start server on port $port"
|
||||
cleanup
|
||||
# shellcheck disable=SC2317
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
@@ -209,8 +210,8 @@ main() {
|
||||
# =============================================================================
|
||||
# Run Benchmark
|
||||
# =============================================================================
|
||||
vllm bench serve --port $PROXY_PORT --seed $(date +%s) \
|
||||
--backend vllm --model $MODEL \
|
||||
vllm bench serve --port "$PROXY_PORT" --seed "$(date +%s)" \
|
||||
--backend vllm --model "$MODEL" \
|
||||
--dataset-name random --random-input-len 7500 --random-output-len 200 \
|
||||
--num-prompts 200 --burstiness 100 --request-rate 2 | tee benchmark.log
|
||||
|
||||
|
||||
+8
-7
@@ -166,10 +166,10 @@ main() {
|
||||
local kv_port=$((21001 + i))
|
||||
|
||||
echo " Prefill server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--host 0.0.0.0 \
|
||||
--port $port \
|
||||
--port "$port" \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--dtype float16 \
|
||||
@@ -194,10 +194,10 @@ main() {
|
||||
local kv_port=$((22001 + i))
|
||||
|
||||
echo " Decode server $((i+1)): GPU $gpu_id, Port $port, KV Port $kv_port"
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve $MODEL \
|
||||
CUDA_VISIBLE_DEVICES=$gpu_id vllm serve "$MODEL" \
|
||||
--enforce-eager \
|
||||
--host 0.0.0.0 \
|
||||
--port $port \
|
||||
--port "$port" \
|
||||
--tensor-parallel-size 1 \
|
||||
--seed 1024 \
|
||||
--dtype float16 \
|
||||
@@ -217,9 +217,10 @@ main() {
|
||||
echo ""
|
||||
echo "Waiting for all servers to start..."
|
||||
for port in "${PREFILL_PORT_ARRAY[@]}" "${DECODE_PORT_ARRAY[@]}"; do
|
||||
if ! wait_for_server $port; then
|
||||
if ! wait_for_server "$port"; then
|
||||
echo "Failed to start server on port $port"
|
||||
cleanup
|
||||
# shellcheck disable=SC2317
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
@@ -231,8 +232,8 @@ main() {
|
||||
# Run Benchmark
|
||||
# =============================================================================
|
||||
cd ../../../benchmarks/
|
||||
vllm bench serve --port 10001 --seed $(date +%s) \
|
||||
--model $MODEL \
|
||||
vllm bench serve --port 10001 --seed "$(date +%s)" \
|
||||
--model "$MODEL" \
|
||||
--dataset-name random --random-input-len 7500 --random-output-len 200 \
|
||||
--num-prompts 200 --burstiness 100 --request-rate 2 | tee benchmark.log
|
||||
|
||||
|
||||
@@ -50,8 +50,8 @@ while [[ $# -gt 0 ]]; do
|
||||
done
|
||||
|
||||
vllm bench serve \
|
||||
--model $MODEL_NAME \
|
||||
--host $HOST \
|
||||
--port $PORT \
|
||||
--num-prompts $NUM_PROMPTS \
|
||||
--request-rate $REQUEST_RATE
|
||||
--model "$MODEL_NAME" \
|
||||
--host "$HOST" \
|
||||
--port "$PORT" \
|
||||
--num-prompts "$NUM_PROMPTS" \
|
||||
--request-rate "$REQUEST_RATE"
|
||||
|
||||
@@ -57,15 +57,15 @@ echo "Starting vLLM server for $MODEL_NAME with data parallel size: $DATA_PARALL
|
||||
export RAY_DEDUP_LOGS=0
|
||||
export VLLM_USE_DEEP_GEMM=1
|
||||
|
||||
vllm serve $MODEL_NAME \
|
||||
--data-parallel-size $DATA_PARALLEL_SIZE \
|
||||
--data-parallel-size-local $DATA_PARALLEL_SIZE \
|
||||
vllm serve "$MODEL_NAME" \
|
||||
--data-parallel-size "$DATA_PARALLEL_SIZE" \
|
||||
--data-parallel-size-local "$DATA_PARALLEL_SIZE" \
|
||||
--data-parallel-backend ray \
|
||||
--enforce-eager \
|
||||
--enable-expert-parallel \
|
||||
--enable-eplb \
|
||||
--all2all-backend pplx \
|
||||
--num-redundant-experts $REDUNDANT_EXPERTS \
|
||||
--num-redundant-experts "$REDUNDANT_EXPERTS" \
|
||||
--trust-remote-code \
|
||||
--host $HOST \
|
||||
--port $PORT
|
||||
--host "$HOST" \
|
||||
--port "$PORT"
|
||||
|
||||
@@ -57,8 +57,7 @@ case "$subcommand" in
|
||||
|
||||
# Retry until the worker node connects to the head node or the timeout expires.
|
||||
for (( i=0; i < $ray_init_timeout; i+=5 )); do
|
||||
ray start --address=$ray_address:$ray_port --block "${start_params[@]}"
|
||||
if [ $? -eq 0 ]; then
|
||||
if ray start --address="$ray_address":"$ray_port" --block "${start_params[@]}"; then
|
||||
echo "Worker: Ray runtime started with head address $ray_address:$ray_port"
|
||||
exit 0
|
||||
fi
|
||||
@@ -95,12 +94,12 @@ case "$subcommand" in
|
||||
fi
|
||||
|
||||
# Start the Ray head node.
|
||||
ray start --head --port=$ray_port "${start_params[@]}"
|
||||
ray start --head --port="$ray_port" "${start_params[@]}"
|
||||
|
||||
# Poll Ray until every worker node is active.
|
||||
for (( i=0; i < $ray_init_timeout; i+=5 )); do
|
||||
active_nodes=`python3 -c 'import ray; ray.init(); print(sum(node["Alive"] for node in ray.nodes()))'`
|
||||
if [ $active_nodes -eq $ray_cluster_size ]; then
|
||||
active_nodes=$(python3 -c 'import ray; ray.init(); print(sum(node["Alive"] for node in ray.nodes()))')
|
||||
if [ "$active_nodes" -eq "$ray_cluster_size" ]; then
|
||||
echo "All ray workers are active and the ray cluster is initialized successfully."
|
||||
exit 0
|
||||
fi
|
||||
|
||||
@@ -22,11 +22,10 @@ check_hf_token() {
|
||||
|
||||
check_num_gpus() {
|
||||
# can you check if the number of GPUs are >=2 via nvidia-smi/rocm-smi?
|
||||
which rocm-smi > /dev/null 2>&1
|
||||
if [ $? -ne 0 ]; then
|
||||
if ! which rocm-smi > /dev/null 2>&1; then
|
||||
num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
|
||||
else
|
||||
num_gpus=$(rocm-smi --showid | grep Instinct | wc -l)
|
||||
num_gpus=$(rocm-smi --showid | grep -c Instinct)
|
||||
fi
|
||||
|
||||
if [ "$num_gpus" -lt 2 ]; then
|
||||
@@ -39,8 +38,7 @@ check_num_gpus() {
|
||||
|
||||
ensure_python_library_installed() {
|
||||
echo "Checking if $1 is installed..."
|
||||
python3 -c "import $1" > /dev/null 2>&1
|
||||
if [ $? -ne 0 ]; then
|
||||
if ! python3 -c "import $1" > /dev/null 2>&1; then
|
||||
if [ "$1" == "nixl" ]; then
|
||||
echo "$1 is not installed. Please refer to https://github.com/ai-dynamo/nixl for installation."
|
||||
else
|
||||
@@ -102,12 +100,12 @@ main() {
|
||||
bash disagg_vllm_launcher.sh prefiller \
|
||||
> >(tee prefiller.log) 2>&1 &
|
||||
prefiller_pid=$!
|
||||
PIDS+=($prefiller_pid)
|
||||
PIDS+=("$prefiller_pid")
|
||||
|
||||
bash disagg_vllm_launcher.sh decoder \
|
||||
> >(tee decoder.log) 2>&1 &
|
||||
decoder_pid=$!
|
||||
PIDS+=($decoder_pid)
|
||||
PIDS+=("$decoder_pid")
|
||||
|
||||
python3 disagg_proxy_server.py \
|
||||
--host localhost \
|
||||
@@ -118,7 +116,7 @@ main() {
|
||||
--decoder-port 8200 \
|
||||
> >(tee proxy.log) 2>&1 &
|
||||
proxy_pid=$!
|
||||
PIDS+=($proxy_pid)
|
||||
PIDS+=("$proxy_pid")
|
||||
|
||||
wait_for_server 8100
|
||||
wait_for_server 8200
|
||||
@@ -128,7 +126,7 @@ main() {
|
||||
|
||||
# begin benchmark
|
||||
cd ../../../../benchmarks/
|
||||
vllm bench serve --port 9000 --seed $(date +%s) \
|
||||
vllm bench serve --port 9000 --seed "$(date +%s)" \
|
||||
--model meta-llama/Llama-3.1-8B-Instruct \
|
||||
--dataset-name random --random-input-len 7500 --random-output-len 200 \
|
||||
--num-prompts 200 --burstiness 100 --request-rate 3.6 | tee benchmark.log
|
||||
|
||||
@@ -34,7 +34,7 @@ if [[ $1 == "prefiller" ]]; then
|
||||
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
CUDA_VISIBLE_DEVICES=0 \
|
||||
vllm serve $MODEL \
|
||||
vllm serve "$MODEL" \
|
||||
--port 8100 \
|
||||
--enforce-eager \
|
||||
--kv-transfer-config \
|
||||
@@ -51,7 +51,7 @@ elif [[ $1 == "decoder" ]]; then
|
||||
VLLM_ENABLE_V1_MULTIPROCESSING=1 \
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
CUDA_VISIBLE_DEVICES=1 \
|
||||
vllm serve $MODEL \
|
||||
vllm serve "$MODEL" \
|
||||
--port 8200 \
|
||||
--enforce-eager \
|
||||
--kv-transfer-config \
|
||||
|
||||
@@ -103,7 +103,7 @@ vllm serve "$MODEL_NAME" \
|
||||
--tensor-parallel-size "$GPU_COUNT" \
|
||||
--enforce-eager \
|
||||
--pooler-config "$POOLER_CONFIG" \
|
||||
--served-model-name ${MODEL_CODE} \
|
||||
--served-model-name "${MODEL_CODE}" \
|
||||
--api-key "$API_KEY" \
|
||||
--trust-remote-code \
|
||||
--port "$PORT" \
|
||||
|
||||
@@ -20,13 +20,15 @@ def main():
|
||||
torch.set_default_dtype(torch.float16)
|
||||
image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" # noqa: E501
|
||||
|
||||
img_prompt = dict(
|
||||
img_data = dict(
|
||||
data=image_url,
|
||||
data_format="url",
|
||||
image_format="tiff",
|
||||
out_data_format="b64_json",
|
||||
)
|
||||
|
||||
prompt = dict(data=img_data)
|
||||
|
||||
llm = LLM(
|
||||
model="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
|
||||
skip_tokenizer_init=True,
|
||||
@@ -41,7 +43,7 @@ def main():
|
||||
enable_mm_embeds=True,
|
||||
)
|
||||
|
||||
pooler_output = llm.encode(img_prompt, pooling_task="plugin")
|
||||
pooler_output = llm.encode(prompt, pooling_task="plugin")
|
||||
output = pooler_output[0].outputs
|
||||
|
||||
print(output)
|
||||
|
||||
@@ -105,6 +105,10 @@ plugins:
|
||||
- https://numpy.org/doc/stable/objects.inv
|
||||
- https://pytorch.org/docs/stable/objects.inv
|
||||
- https://psutil.readthedocs.io/en/stable/objects.inv
|
||||
- redirects:
|
||||
redirect_maps:
|
||||
features/spec_decode/README.md: features/speculative_decoding/README.md
|
||||
features/spec_decode/speculators.md: features/speculative_decoding/speculators.md
|
||||
|
||||
markdown_extensions:
|
||||
- attr_list
|
||||
|
||||
@@ -113,7 +113,6 @@ markers = [
|
||||
"cpu_test: mark test as CPU-only test",
|
||||
"split: run this test as part of a split",
|
||||
"distributed: run this test only in distributed GPU tests",
|
||||
"skip_v1: do not run this test with v1",
|
||||
"optional: optional tests that are automatically skipped, include --optional to run them",
|
||||
]
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ mkdocs-awesome-nav
|
||||
mkdocs-glightbox
|
||||
mkdocs-git-revision-date-localized-plugin
|
||||
mkdocs-minify-plugin
|
||||
mkdocs-redirects
|
||||
regex
|
||||
ruff
|
||||
pydantic
|
||||
|
||||
@@ -1,2 +1,3 @@
|
||||
lmcache >= 0.3.9
|
||||
nixl >= 0.7.1 # Required for disaggregated prefill
|
||||
mooncake-transfer-engine >= 0.3.8
|
||||
|
||||
@@ -27,7 +27,7 @@ mistral_common[image,audio] >= 1.9.1 # required for voxtral test
|
||||
num2words # required for smolvlm test
|
||||
opencv-python-headless >= 4.13.0 # required for video test
|
||||
datamodel_code_generator # required for minicpm3 test
|
||||
lm-eval[api]>=0.4.9.2 # required for model evaluation test
|
||||
lm-eval[api]>=0.4.11 # required for model evaluation test
|
||||
mteb>=1.38.11, <2 # required for mteb test
|
||||
transformers==4.57.5
|
||||
tokenizers==0.22.0
|
||||
|
||||
@@ -58,7 +58,7 @@ schemathesis==3.39.15
|
||||
# OpenAI schema test
|
||||
|
||||
# Evaluation and benchmarking
|
||||
lm-eval[api]==0.4.9.2
|
||||
lm-eval[api]==0.4.11
|
||||
jiwer==4.0.0
|
||||
|
||||
# Required for multiprocessed tests that use spawn method, Datasets and Evaluate Test
|
||||
@@ -67,8 +67,6 @@ multiprocess==0.70.16
|
||||
# Required for v1/metrics/test_engine_logger_apis.py
|
||||
ray[cgraph,default]>=2.48.0
|
||||
|
||||
# Plugins test
|
||||
terratorch @ git+https://github.com/IBM/terratorch.git@07184fcf91a1324f831ff521dd238d97fe350e3e
|
||||
torchgeo==0.7.0
|
||||
# via terratorch
|
||||
# MTEB Benchmark Test
|
||||
@@ -98,3 +96,9 @@ transformers==4.57.3
|
||||
huggingface-hub==0.36.2
|
||||
# Pin Mistral Common
|
||||
mistral-common[image,audio]==1.9.1
|
||||
# Required for Prithvi tests
|
||||
terratorch==1.2.2
|
||||
# Required for Prithvi tests
|
||||
segmentation-models-pytorch==0.5.0
|
||||
# Required for Prithvi tests
|
||||
imagehash==4.3.2
|
||||
|
||||
@@ -35,7 +35,7 @@ num2words # required for smolvlm test
|
||||
open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
|
||||
opencv-python-headless >= 4.13.0 # required for video test
|
||||
datamodel_code_generator # required for minicpm3 test
|
||||
lm-eval[api]>=0.4.9.2 # required for model evaluation test
|
||||
lm-eval[api]>=0.4.11 # required for model evaluation test
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
transformers==4.57.5
|
||||
tokenizers==0.22.0
|
||||
|
||||
+6
-18
@@ -5,9 +5,7 @@ absl-py==2.1.0
|
||||
# rouge-score
|
||||
# tensorboard
|
||||
accelerate==1.0.1
|
||||
# via
|
||||
# lm-eval
|
||||
# peft
|
||||
# via peft
|
||||
aenum==3.1.16
|
||||
# via lightly
|
||||
affine==2.4.0
|
||||
@@ -138,7 +136,6 @@ colorama==0.4.6
|
||||
# perceptron
|
||||
# sacrebleu
|
||||
# schemathesis
|
||||
# tqdm-multiprocess
|
||||
colorful==0.5.6
|
||||
# via ray
|
||||
colorlog==6.10.1
|
||||
@@ -383,6 +380,7 @@ jinja2==3.1.6
|
||||
# via
|
||||
# datamodel-code-generator
|
||||
# genai-perf
|
||||
# lm-eval
|
||||
# torch
|
||||
jiwer==3.0.5
|
||||
# via -r requirements/test.in
|
||||
@@ -448,7 +446,7 @@ lightning-utilities==0.14.3
|
||||
# torchmetrics
|
||||
llvmlite==0.44.0
|
||||
# via numba
|
||||
lm-eval==0.4.9.2
|
||||
lm-eval==0.4.11
|
||||
# via -r requirements/test.in
|
||||
lxml==5.3.0
|
||||
# via
|
||||
@@ -513,8 +511,6 @@ numba==0.61.2
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# librosa
|
||||
numexpr==2.10.1
|
||||
# via lm-eval
|
||||
numpy==2.2.6
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
@@ -540,11 +536,11 @@ numpy==2.2.6
|
||||
# librosa
|
||||
# lightly
|
||||
# lightly-utils
|
||||
# lm-eval
|
||||
# matplotlib
|
||||
# mistral-common
|
||||
# mteb
|
||||
# numba
|
||||
# numexpr
|
||||
# opencv-python-headless
|
||||
# optuna
|
||||
# pandas
|
||||
@@ -707,9 +703,7 @@ pathvalidate==3.2.1
|
||||
patsy==1.0.1
|
||||
# via statsmodels
|
||||
peft==0.16.0
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# lm-eval
|
||||
# via -r requirements/test.in
|
||||
perceptron==0.1.4
|
||||
# via -r requirements/test.in
|
||||
perf-analyzer==0.1.0
|
||||
@@ -792,8 +786,6 @@ pyasn1==0.6.1
|
||||
# rsa
|
||||
pyasn1-modules==0.4.2
|
||||
# via google-auth
|
||||
pybind11==2.13.6
|
||||
# via lm-eval
|
||||
pycocotools==2.0.8
|
||||
# via terratorch
|
||||
pycountry==24.6.1
|
||||
@@ -1171,7 +1163,6 @@ torch==2.10.0+cu129
|
||||
# kornia
|
||||
# lightly
|
||||
# lightning
|
||||
# lm-eval
|
||||
# mteb
|
||||
# open-clip-torch
|
||||
# peft
|
||||
@@ -1229,15 +1220,11 @@ tqdm==4.67.3
|
||||
# sentence-transformers
|
||||
# tacoreader
|
||||
# terratorch
|
||||
# tqdm-multiprocess
|
||||
# transformers
|
||||
tqdm-multiprocess==0.0.11
|
||||
# via lm-eval
|
||||
transformers==4.57.5
|
||||
# via
|
||||
# -r requirements/test.in
|
||||
# genai-perf
|
||||
# lm-eval
|
||||
# peft
|
||||
# sentence-transformers
|
||||
# transformers-stream-generator
|
||||
@@ -1272,6 +1259,7 @@ typing-extensions==4.15.0
|
||||
# librosa
|
||||
# lightning
|
||||
# lightning-utilities
|
||||
# lm-eval
|
||||
# mistral-common
|
||||
# mteb
|
||||
# opentelemetry-api
|
||||
|
||||
@@ -977,9 +977,9 @@ if _is_cuda():
|
||||
# FA3 requires CUDA 12.3 or later
|
||||
ext_modules.append(CMakeExtension(name="vllm.vllm_flash_attn._vllm_fa3_C"))
|
||||
if envs.VLLM_USE_PRECOMPILED or (
|
||||
CUDA_HOME and get_nvcc_cuda_version() >= Version("12.9")
|
||||
CUDA_HOME and get_nvcc_cuda_version() >= Version("12.8")
|
||||
):
|
||||
# FlashMLA requires CUDA 12.9 or later
|
||||
# FlashMLA requires CUDA 12.8 or later
|
||||
# Optional since this doesn't get built (produce an .so file) when
|
||||
# not targeting a hopper system
|
||||
ext_modules.append(CMakeExtension(name="vllm._flashmla_C", optional=True))
|
||||
|
||||
@@ -229,7 +229,7 @@ def _compare_sp(
|
||||
if chunked_prefill:
|
||||
common_args.append("--enable-chunked-prefill")
|
||||
if eager_mode:
|
||||
common_args.append("--enforce-eager")
|
||||
common_args.append("-cc.cudagraph_mode=none")
|
||||
if runner != "auto":
|
||||
common_args.extend(["--runner", runner])
|
||||
if trust_remote_code:
|
||||
|
||||
@@ -50,10 +50,9 @@ def test_tp1_fp8_fusions(
|
||||
run_e2e_fusion_test,
|
||||
monkeypatch,
|
||||
):
|
||||
if use_deepgemm:
|
||||
# TODO(luka/eliza) DeepGEMM uses different quants, matching not supported
|
||||
if use_deepgemm and is_blackwell():
|
||||
# TODO(luka) DeepGEMM uses different quants, matching not supported
|
||||
# - on Blackwell, uses a special quant fp8, currently not supported
|
||||
# - on Hopper, tma-aligned scales inhibit matching (fix WIP)
|
||||
pytest.skip("DeepGEMM & quant matching not currently supported")
|
||||
|
||||
matches = matches_fn(n_layers)
|
||||
@@ -66,7 +65,6 @@ def test_tp1_fp8_fusions(
|
||||
model_kwargs["hf_overrides"] = hf_overrides(n_layers)
|
||||
model_kwargs["load_format"] = "dummy"
|
||||
model_kwargs["max_model_len"] = 1024
|
||||
|
||||
compilation_config = dict(
|
||||
use_inductor_graph_partition=inductor_graph_partition,
|
||||
custom_ops=custom_ops.split(","),
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user