Compare commits
87
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6a7b85d940 | ||
|
|
9f7ee8921a | ||
|
|
fcd6ac97ed | ||
|
|
95be2a7f22 | ||
|
|
0e60c925cf | ||
|
|
d7ff22204a | ||
|
|
4726250123 | ||
|
|
c0bd8b13da | ||
|
|
caeb887bf6 | ||
|
|
6b3166a7c7 | ||
|
|
25e2e136ef | ||
|
|
6874638bc4 | ||
|
|
e24663c5a9 | ||
|
|
c50e105a88 | ||
|
|
a766b30349 | ||
|
|
1faa8cb73c | ||
|
|
e89a91d927 | ||
|
|
909b147197 | ||
|
|
a88b3be7c4 | ||
|
|
b58a393fd4 | ||
|
|
ec5bf35ce4 | ||
|
|
a49ea5a58f | ||
|
|
522f51d7fb | ||
|
|
47b9319506 | ||
|
|
31b4eff44c | ||
|
|
12204cf89b | ||
|
|
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
|
||||
|
||||
@@ -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 692917b1cda61b93ac9ee2d846ec54e75afe87b1
|
||||
GIT_PROGRESS TRUE
|
||||
CONFIGURE_COMMAND ""
|
||||
BUILD_COMMAND ""
|
||||
|
||||
@@ -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);
|
||||
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(","),
|
||||
|
||||
@@ -7,7 +7,6 @@ from vllm.entrypoints.llm import LLM
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
|
||||
@pytest.mark.skip_v1
|
||||
@pytest.mark.parametrize("model", ["distilbert/distilgpt2"])
|
||||
def test_computed_prefix_blocks(model: str):
|
||||
# This test checks if the engine generates completions both with and
|
||||
|
||||
@@ -145,6 +145,7 @@ async def test_request_cancellation(server: RemoteOpenAIServer):
|
||||
model=MODEL_NAME,
|
||||
max_tokens=10000,
|
||||
extra_body={"min_tokens": 10000},
|
||||
temperature=0.0,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
@@ -163,7 +164,7 @@ async def test_request_cancellation(server: RemoteOpenAIServer):
|
||||
# be able to respond to this one within the timeout
|
||||
client = server.get_async_client(timeout=5)
|
||||
response = await client.chat.completions.create(
|
||||
messages=chat_input, model=MODEL_NAME, max_tokens=10
|
||||
messages=chat_input, model=MODEL_NAME, max_tokens=10, temperature=0.0
|
||||
)
|
||||
|
||||
assert len(response.choices) == 1
|
||||
|
||||
@@ -17,6 +17,7 @@ from transformers import AutoTokenizer
|
||||
from tests.conftest import LocalAssetServer
|
||||
from tests.utils import RemoteOpenAIServer
|
||||
from vllm import version
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
|
||||
MODELS = {
|
||||
"text": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
|
||||
@@ -315,14 +316,26 @@ async def test_abort_metrics_reset(
|
||||
client.completions.create(
|
||||
model=model_name,
|
||||
prompt=prompt_ids,
|
||||
max_tokens=100, # Long generation to give time to abort
|
||||
max_tokens=500, # Long generation to give time to abort
|
||||
temperature=0.0,
|
||||
)
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
# Wait a bit for requests to start processing
|
||||
await asyncio.sleep(0.5)
|
||||
# Poll until we see running requests rather than using a fixed sleep,
|
||||
# since generation speed varies across hardware.
|
||||
try:
|
||||
await _poll_until(
|
||||
lambda: _get_running_metrics_from_api(server)[0] > 0,
|
||||
timeout=10.0,
|
||||
interval=0.1,
|
||||
description="running_requests > 0",
|
||||
)
|
||||
except TimeoutError:
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
pytest.fail("Requests never appeared as running in metrics")
|
||||
|
||||
# Check that we have running requests
|
||||
running_requests, waiting_requests, kv_cache_usage = _get_running_metrics_from_api(
|
||||
@@ -336,13 +349,15 @@ async def test_abort_metrics_reset(
|
||||
# Cancel all tasks to abort the requests
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
# Wait for cancellations to be processed
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
# Check that metrics have reset to zero
|
||||
response = requests.get(server.url_for("metrics"))
|
||||
assert response.status_code == HTTPStatus.OK
|
||||
# Poll until metrics reset rather than using a fixed sleep
|
||||
await _poll_until(
|
||||
lambda: _get_running_metrics_from_api(server) == (0, 0, 0),
|
||||
timeout=10.0,
|
||||
interval=0.2,
|
||||
description="gauge metrics back to zero",
|
||||
)
|
||||
|
||||
# Verify running and waiting requests counts and KV cache usage are zero
|
||||
running_requests_after, waiting_requests_after, kv_cache_usage_after = (
|
||||
@@ -360,6 +375,18 @@ async def test_abort_metrics_reset(
|
||||
)
|
||||
|
||||
|
||||
async def _poll_until(
|
||||
predicate, *, timeout: float, interval: float = 0.5, description: str = "condition"
|
||||
):
|
||||
"""Poll until predicate() returns True, or raise TimeoutError."""
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
if predicate():
|
||||
return
|
||||
await asyncio.sleep(interval)
|
||||
raise TimeoutError(f"Timed out after {timeout}s waiting for: {description}")
|
||||
|
||||
|
||||
def _get_running_metrics_from_api(server: RemoteOpenAIServer):
|
||||
"""Return (running_count, waiting_count, kv_cache_usage)"""
|
||||
|
||||
@@ -399,7 +426,7 @@ def test_metrics_exist_run_batch():
|
||||
input_batch = """{"custom_id": "request-0", "method": "POST", "url": "/v1/embeddings", "body": {"model": "intfloat/multilingual-e5-small", "input": "You are a helpful assistant."}}""" # noqa: E501
|
||||
|
||||
base_url = "0.0.0.0"
|
||||
port = "8001"
|
||||
port = str(get_open_port())
|
||||
server_url = f"http://{base_url}:{port}"
|
||||
|
||||
with (
|
||||
@@ -427,17 +454,32 @@ def test_metrics_exist_run_batch():
|
||||
],
|
||||
)
|
||||
|
||||
def is_server_up(url):
|
||||
try:
|
||||
|
||||
def is_server_up(url):
|
||||
try:
|
||||
response = requests.get(url)
|
||||
return response.status_code == 200
|
||||
except requests.ConnectionError:
|
||||
return False
|
||||
|
||||
start = time.time()
|
||||
timeout = 120
|
||||
while not is_server_up(server_url):
|
||||
if proc.poll() is not None:
|
||||
pytest.fail(
|
||||
f"Batch process exited early with returncode={proc.returncode}"
|
||||
)
|
||||
if time.time() - start > timeout:
|
||||
pytest.fail("Batch server did not start within timeout")
|
||||
time.sleep(1)
|
||||
|
||||
response = requests.get(server_url + "/metrics")
|
||||
assert response.status_code == HTTPStatus.OK
|
||||
finally:
|
||||
proc.terminate()
|
||||
try:
|
||||
response = requests.get(url)
|
||||
return response.status_code == 200
|
||||
except requests.ConnectionError:
|
||||
return False
|
||||
|
||||
while not is_server_up(server_url):
|
||||
time.sleep(1)
|
||||
|
||||
response = requests.get(server_url + "/metrics")
|
||||
assert response.status_code == HTTPStatus.OK
|
||||
|
||||
proc.wait()
|
||||
proc.wait(timeout=15)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.wait(timeout=5)
|
||||
|
||||
@@ -195,18 +195,15 @@ def test_chat_batch_failure_cleanup(llm_for_failure_test):
|
||||
valid_msg = [{"role": "user", "content": "Hello"}]
|
||||
long_text = "This is a very long text to test the error " * 50
|
||||
invalid_msg = [{"role": "user", "content": long_text}]
|
||||
batch_1 = [
|
||||
valid_msg,
|
||||
valid_msg,
|
||||
invalid_msg,
|
||||
]
|
||||
batch_2 = [
|
||||
valid_msg,
|
||||
valid_msg,
|
||||
]
|
||||
|
||||
batch_1 = [valid_msg, valid_msg, invalid_msg]
|
||||
batch_2 = [valid_msg, valid_msg]
|
||||
sampling_params = SamplingParams(temperature=0, max_tokens=10)
|
||||
|
||||
with pytest.raises(ValueError, match="context length is only"):
|
||||
llm.chat(batch_1, sampling_params=sampling_params)
|
||||
assert llm.llm_engine.get_num_unfinished_requests() == 0
|
||||
|
||||
outputs_2 = llm.chat(batch_2, sampling_params=sampling_params)
|
||||
assert len(outputs_2) == len(batch_2)
|
||||
assert llm.llm_engine.get_num_unfinished_requests() == 0
|
||||
|
||||
@@ -6,7 +6,6 @@ import time
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
import requests
|
||||
from openai import BadRequestError, NotFoundError, OpenAI
|
||||
from openai_harmony import (
|
||||
Message,
|
||||
@@ -513,11 +512,9 @@ async def test_code_interpreter(client: OpenAI, model_name: str):
|
||||
|
||||
|
||||
def get_weather(latitude, longitude):
|
||||
response = requests.get(
|
||||
f"https://api.open-meteo.com/v1/forecast?latitude={latitude}&longitude={longitude}¤t=temperature_2m,wind_speed_10m&hourly=temperature_2m,relative_humidity_2m,wind_speed_10m" # noqa
|
||||
)
|
||||
data = response.json()
|
||||
return data["current"]["temperature_2m"]
|
||||
# Return a static temperature value to avoid flaky SSL/network errors
|
||||
# from calling the external api.open-meteo.com API in CI.
|
||||
return 15.0
|
||||
|
||||
|
||||
def get_place_to_travel():
|
||||
|
||||
@@ -172,19 +172,26 @@ async def test_mcp_tool_call(client: OpenAI, model_name: str):
|
||||
|
||||
assert response is not None
|
||||
assert response.status == "completed"
|
||||
assert response.output[0].type == "reasoning"
|
||||
assert response.output[1].type == "mcp_call"
|
||||
assert type(response.output[1].arguments) is str
|
||||
assert type(response.output[1].output) is str
|
||||
assert response.output[2].type == "reasoning"
|
||||
# make sure the correct math is in the final output
|
||||
assert response.output[3].type == "message"
|
||||
assert any(s in response.output[3].content[0].text for s in ("56088", "56,088"))
|
||||
|
||||
# test raw input_messages / output_messages
|
||||
# The model may produce multiple reasoning/mcp_call rounds before the
|
||||
# final message, so validate structurally rather than by exact index.
|
||||
output_types = [o.type for o in response.output]
|
||||
assert "reasoning" in output_types
|
||||
mcp_calls = [o for o in response.output if o.type == "mcp_call"]
|
||||
assert len(mcp_calls) >= 1
|
||||
assert type(mcp_calls[0].arguments) is str
|
||||
assert type(mcp_calls[0].output) is str
|
||||
|
||||
# The final output should be a message containing the correct answer
|
||||
assert response.output[-1].type == "message"
|
||||
assert any(s in response.output[-1].content[0].text for s in ("56088", "56,088"))
|
||||
|
||||
# Test raw input_messages / output_messages
|
||||
assert len(response.input_messages) == 1
|
||||
assert len(response.output_messages) == 3
|
||||
assert any(s in response.output_messages[2]["message"] for s in ("56088", "56,088"))
|
||||
assert len(response.output_messages) >= 3
|
||||
assert any(
|
||||
s in response.output_messages[-1]["message"] for s in ("56088", "56,088")
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -151,6 +151,7 @@ def test_openapi_stateless(case: schemathesis.Case):
|
||||
# requires a longer timeout
|
||||
("POST", "/v1/chat/completions"): LONG_TIMEOUT_SECONDS,
|
||||
("POST", "/v1/completions"): LONG_TIMEOUT_SECONDS,
|
||||
("POST", "/v1/messages"): LONG_TIMEOUT_SECONDS,
|
||||
}.get(key, DEFAULT_TIMEOUT_SECONDS)
|
||||
|
||||
# No need to verify SSL certificate for localhost
|
||||
|
||||
@@ -526,6 +526,7 @@ class MockModelConfig:
|
||||
allowed_media_domains: list[str] | None = None
|
||||
encoder_config = None
|
||||
generation_config: str = "auto"
|
||||
override_generation_config: dict[str, Any] = field(default_factory=dict)
|
||||
media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
|
||||
skip_tokenizer_init: bool = False
|
||||
is_encoder_decoder: bool = False
|
||||
@@ -651,12 +652,10 @@ async def test_serving_chat_should_set_correct_max_tokens():
|
||||
|
||||
assert mock_engine.generate.call_args.args[1].max_tokens == 10
|
||||
|
||||
# Setting server's max_tokens in the generation_config.json
|
||||
# lower than context_window - prompt_tokens
|
||||
# Model author's generation_config.json sets max_tokens (auto, no override)
|
||||
# — should act as fallback only, not ceiling
|
||||
mock_model_config = MockModelConfig()
|
||||
mock_model_config.diff_sampling_param = {
|
||||
"max_tokens": 10 # Setting server-side max_tokens limit
|
||||
}
|
||||
mock_model_config.diff_sampling_param = {"max_tokens": 10}
|
||||
|
||||
# Reinitialize the engine with new settings
|
||||
mock_engine = MagicMock(spec=AsyncLLM)
|
||||
@@ -680,7 +679,50 @@ async def test_serving_chat_should_set_correct_max_tokens():
|
||||
|
||||
assert mock_engine.generate.call_args.args[1].max_tokens == 10
|
||||
|
||||
# Test Case 2: Request's max_tokens set higher than server accepts
|
||||
# Test Case 2: Request's max_tokens set higher than generation_config
|
||||
# default so request-provided max_tokens takes precedence
|
||||
req.max_tokens = 15
|
||||
|
||||
with suppress(Exception):
|
||||
await serving_chat.create_chat_completion(req)
|
||||
|
||||
assert mock_engine.generate.call_args.args[1].max_tokens == 15
|
||||
|
||||
# Test Case 3: Request's max_tokens set lower than server accepts
|
||||
req.max_tokens = 5
|
||||
|
||||
with suppress(Exception):
|
||||
await serving_chat.create_chat_completion(req)
|
||||
|
||||
assert mock_engine.generate.call_args.args[1].max_tokens == 5
|
||||
|
||||
# User explicitly sets max_tokens via --override-generation-config
|
||||
# — should act as a ceiling
|
||||
mock_model_config = MockModelConfig()
|
||||
mock_model_config.diff_sampling_param = {"max_tokens": 10}
|
||||
mock_model_config.override_generation_config = {"max_new_tokens": 10}
|
||||
|
||||
mock_engine = MagicMock(spec=AsyncLLM)
|
||||
mock_engine.errored = False
|
||||
mock_engine.model_config = mock_model_config
|
||||
mock_engine.input_processor = MagicMock()
|
||||
mock_engine.io_processor = MagicMock()
|
||||
mock_engine.renderer = _build_renderer(mock_engine.model_config)
|
||||
|
||||
serving_chat = _build_serving_chat(mock_engine)
|
||||
|
||||
# Test Case 3.1: No max_tokens — uses override as default
|
||||
req = ChatCompletionRequest(
|
||||
model=MODEL_NAME,
|
||||
messages=[{"role": "user", "content": "what is 1+1?"}],
|
||||
)
|
||||
|
||||
with suppress(Exception):
|
||||
await serving_chat.create_chat_completion(req)
|
||||
|
||||
assert mock_engine.generate.call_args.args[1].max_tokens == 10
|
||||
|
||||
# Test Case 3.2: Request max_tokens higher — capped by user ceiling from override
|
||||
req.max_tokens = 15
|
||||
|
||||
with suppress(Exception):
|
||||
@@ -688,7 +730,7 @@ async def test_serving_chat_should_set_correct_max_tokens():
|
||||
|
||||
assert mock_engine.generate.call_args.args[1].max_tokens == 10
|
||||
|
||||
# Test Case 3: Request's max_tokens set lower than server accepts
|
||||
# Test Case 3.3: Request max_tokens lower — respected
|
||||
req.max_tokens = 5
|
||||
|
||||
with suppress(Exception):
|
||||
@@ -699,9 +741,7 @@ async def test_serving_chat_should_set_correct_max_tokens():
|
||||
# Setting server's max_tokens in the generation_config.json
|
||||
# higher than context_window - prompt_tokens
|
||||
mock_model_config = MockModelConfig()
|
||||
mock_model_config.diff_sampling_param = {
|
||||
"max_tokens": 200 # Setting server-side max_tokens limit
|
||||
}
|
||||
mock_model_config.diff_sampling_param = {"max_tokens": 200}
|
||||
|
||||
# Reinitialize the engine with new settings
|
||||
mock_engine = MagicMock(spec=AsyncLLM)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from vllm.entrypoints.utils import sanitize_message
|
||||
|
||||
from vllm.entrypoints.utils import get_max_tokens, sanitize_message
|
||||
|
||||
|
||||
def test_sanitize_message():
|
||||
@@ -8,3 +9,74 @@ def test_sanitize_message():
|
||||
sanitize_message("<_io.BytesIO object at 0x7a95e299e750>")
|
||||
== "<_io.BytesIO object>"
|
||||
)
|
||||
|
||||
|
||||
class TestGetMaxTokens:
|
||||
"""Tests for get_max_tokens() to ensure generation_config's max_tokens
|
||||
acts as a default when from model author, and as a ceiling when
|
||||
explicitly set by the user."""
|
||||
|
||||
def test_default_sampling_params_used_when_no_request_max_tokens(self):
|
||||
"""When user doesn't specify max_tokens, generation_config default
|
||||
should apply."""
|
||||
result = get_max_tokens(
|
||||
max_model_len=24000,
|
||||
max_tokens=None,
|
||||
input_length=100,
|
||||
default_sampling_params={"max_tokens": 2048},
|
||||
)
|
||||
assert result == 2048
|
||||
|
||||
def test_request_max_tokens_not_capped_by_default_sampling_params(self):
|
||||
"""When user specifies max_tokens in request, model author's
|
||||
generation_config max_tokens must NOT cap it (fixes #34005)."""
|
||||
result = get_max_tokens(
|
||||
max_model_len=24000,
|
||||
max_tokens=5000,
|
||||
input_length=100,
|
||||
default_sampling_params={"max_tokens": 2048},
|
||||
)
|
||||
assert result == 5000
|
||||
|
||||
def test_override_max_tokens_caps_request(self):
|
||||
"""When user explicitly sets max_tokens, it acts as a ceiling."""
|
||||
result = get_max_tokens(
|
||||
max_model_len=24000,
|
||||
max_tokens=5000,
|
||||
input_length=100,
|
||||
default_sampling_params={"max_tokens": 2048},
|
||||
override_max_tokens=2048,
|
||||
)
|
||||
assert result == 2048
|
||||
|
||||
def test_override_max_tokens_used_as_default(self):
|
||||
"""When no request max_tokens, override still applies as default."""
|
||||
result = get_max_tokens(
|
||||
max_model_len=24000,
|
||||
max_tokens=None,
|
||||
input_length=100,
|
||||
default_sampling_params={"max_tokens": 2048},
|
||||
override_max_tokens=2048,
|
||||
)
|
||||
assert result == 2048
|
||||
|
||||
def test_max_model_len_still_caps_output(self):
|
||||
"""max_model_len - input_length is always the hard ceiling."""
|
||||
result = get_max_tokens(
|
||||
max_model_len=3000,
|
||||
max_tokens=5000,
|
||||
input_length=100,
|
||||
default_sampling_params={"max_tokens": 2048},
|
||||
)
|
||||
assert result == 2900 # 3000 - 100
|
||||
|
||||
def test_request_max_tokens_smaller_than_default(self):
|
||||
"""When user explicitly requests fewer tokens than gen_config default,
|
||||
that should be respected."""
|
||||
result = get_max_tokens(
|
||||
max_model_len=24000,
|
||||
max_tokens=512,
|
||||
input_length=100,
|
||||
default_sampling_params={"max_tokens": 2048},
|
||||
)
|
||||
assert result == 512
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
|
||||
accuracy_threshold: 0.29
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP8: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "latency"
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4"
|
||||
accuracy_threshold: 0.29
|
||||
num_questions: 1319
|
||||
num_fewshot: 5
|
||||
server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
|
||||
env:
|
||||
VLLM_USE_FLASHINFER_MOE_FP4: "1"
|
||||
VLLM_FLASHINFER_MOE_BACKEND: "throughput"
|
||||
@@ -13,3 +13,5 @@ Llama-4-Scout-BF16-fi-cutlass.yaml
|
||||
Llama-4-Scout-BF16-triton.yaml
|
||||
Mixtral-8x7B-BF16-fi-cutlass.yaml
|
||||
Mixtral-8x7B-BF16-triton.yaml
|
||||
Nemotron-Nano-30B-Fp8-ModelOpt-fi-trtllm.yaml
|
||||
Nemotron-Nano-30B-NvFp4-ModelOpt-fi-cutlass.yaml
|
||||
|
||||
@@ -17,7 +17,7 @@ from vllm.platforms import current_platform
|
||||
from vllm.platforms.cpu import CpuPlatform
|
||||
from vllm.platforms.cuda import CudaPlatform
|
||||
from vllm.platforms.rocm import RocmPlatform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.utils.torch_utils import set_default_torch_dtype, set_random_seed
|
||||
from vllm.v1.attention.backends.registry import AttentionBackendEnum
|
||||
from vllm.v1.attention.selector import _cached_get_attn_backend
|
||||
|
||||
@@ -71,6 +71,15 @@ def test_mha_attn_platform(default_vllm_config, device: str):
|
||||
attn = MMEncoderAttention(16, 72, scale=1)
|
||||
assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN
|
||||
|
||||
# Test CUDA with head_size=72 (not divisible by 32)
|
||||
# - should use vLLM's FlashAttention
|
||||
with (
|
||||
patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),
|
||||
set_default_torch_dtype(torch.float32),
|
||||
):
|
||||
attn = MMEncoderAttention(16, 72, scale=1)
|
||||
assert attn.attn_backend == AttentionBackendEnum.TRITON_ATTN
|
||||
|
||||
|
||||
def ref_attention(
|
||||
query: torch.Tensor,
|
||||
@@ -153,7 +162,12 @@ def test_mha_attn_forward(
|
||||
v,
|
||||
scale=scale,
|
||||
).reshape(batch_size, seq_len, num_heads * head_size)
|
||||
torch.testing.assert_close(output, ref_output)
|
||||
tol_kwargs = (
|
||||
dict(rtol=1e-3, atol=1e-3)
|
||||
if attn.attn_backend == AttentionBackendEnum.TRITON_ATTN
|
||||
else {}
|
||||
)
|
||||
torch.testing.assert_close(output, ref_output, **tol_kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("var_seq_len", VAR_SEQ_LENS)
|
||||
|
||||
@@ -10,7 +10,7 @@ from vllm.utils.math_utils import next_power_of_2
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.attention.ops.triton_unified_attention import unified_attention
|
||||
|
||||
NUM_HEADS = [(4, 4), (8, 2)]
|
||||
NUM_HEADS = [(4, 4), (8, 2), (5, 1)]
|
||||
HEAD_SIZES = [128, 256]
|
||||
BLOCK_SIZES = [16]
|
||||
|
||||
@@ -20,6 +20,8 @@ QDTYPES = (
|
||||
if not current_platform.is_rocm()
|
||||
else [None, torch.float8_e4m3fnuz]
|
||||
)
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
# one value large enough to test overflow in index calculation.
|
||||
# one value small enough to test the schema op check
|
||||
NUM_BLOCKS = [32768, 2048]
|
||||
@@ -217,3 +219,127 @@ def test_triton_unified_attn(
|
||||
torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output - ref_output))}",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens",
|
||||
[
|
||||
[(1, 1328), (5, 18), (129, 463)],
|
||||
[(1, 523), (1, 37), (1, 2011)],
|
||||
[(1, 1)] * 533,
|
||||
[(533, 533)] * 533,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("num_heads", NUM_HEADS)
|
||||
@pytest.mark.parametrize("head_size", HEAD_SIZES)
|
||||
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
|
||||
@pytest.mark.parametrize("sliding_window", [None, 64, 128, 256])
|
||||
@pytest.mark.parametrize("soft_cap", [None, 50.0])
|
||||
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
|
||||
@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
|
||||
@torch.inference_mode()
|
||||
def test_triton_unified_attn_fp16_input_fp8_output(
|
||||
seq_lens: list[tuple[int, int]],
|
||||
num_heads: tuple[int, int],
|
||||
head_size: int,
|
||||
sliding_window: int | None,
|
||||
block_size: int,
|
||||
soft_cap: float | None,
|
||||
num_blocks: int,
|
||||
seq_threshold_3D: int,
|
||||
) -> None:
|
||||
"""Test with fp16 input and fp8 output using output_scale."""
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
set_random_seed(0)
|
||||
num_seqs = len(seq_lens)
|
||||
query_lens = [x[0] for x in seq_lens]
|
||||
kv_lens = [x[1] for x in seq_lens]
|
||||
num_query_heads = num_heads[0]
|
||||
num_kv_heads = num_heads[1]
|
||||
assert num_query_heads % num_kv_heads == 0
|
||||
max_query_len = max(query_lens)
|
||||
max_kv_len = max(kv_lens)
|
||||
window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
|
||||
scale = head_size**-0.5
|
||||
|
||||
dtype = torch.float16
|
||||
query = torch.randn(sum(query_lens), num_query_heads, head_size, dtype=dtype)
|
||||
key_cache = torch.randn(
|
||||
num_blocks, block_size, num_kv_heads, head_size, dtype=dtype
|
||||
)
|
||||
value_cache = torch.randn_like(key_cache)
|
||||
cu_query_lens = torch.tensor([0] + query_lens, dtype=torch.int32).cumsum(
|
||||
dim=0, dtype=torch.int32
|
||||
)
|
||||
kv_lens_tensor = torch.tensor(kv_lens, dtype=torch.int32)
|
||||
|
||||
max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
|
||||
block_tables = torch.randint(
|
||||
0, num_blocks, (num_seqs, max_num_blocks_per_seq), dtype=torch.int32
|
||||
)
|
||||
|
||||
output = torch.empty(sum(query_lens), num_query_heads, head_size, dtype=FP8_DTYPE)
|
||||
|
||||
output_scale = torch.tensor(0.5, dtype=torch.float32)
|
||||
|
||||
num_par_softmax_segments = 16
|
||||
head_size_padded = next_power_of_2(head_size)
|
||||
softmax_segm_output = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_max = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
softmax_segm_expsum = torch.empty(
|
||||
(seq_threshold_3D, num_query_heads, num_par_softmax_segments),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
unified_attention(
|
||||
q=query,
|
||||
k=key_cache,
|
||||
v=value_cache,
|
||||
out=output,
|
||||
cu_seqlens_q=cu_query_lens,
|
||||
seqused_k=kv_lens_tensor,
|
||||
max_seqlen_q=max_query_len,
|
||||
max_seqlen_k=max_kv_len,
|
||||
softmax_scale=scale,
|
||||
causal=True,
|
||||
window_size=window_size,
|
||||
block_table=block_tables,
|
||||
softcap=soft_cap if soft_cap is not None else 0,
|
||||
q_descale=None,
|
||||
k_descale=None,
|
||||
v_descale=None,
|
||||
output_scale=output_scale,
|
||||
seq_threshold_3D=seq_threshold_3D,
|
||||
num_par_softmax_segments=num_par_softmax_segments,
|
||||
softmax_segm_output=softmax_segm_output,
|
||||
softmax_segm_max=softmax_segm_max,
|
||||
softmax_segm_expsum=softmax_segm_expsum,
|
||||
)
|
||||
|
||||
ref_output = ref_paged_attn(
|
||||
query=query,
|
||||
key_cache=key_cache,
|
||||
value_cache=value_cache,
|
||||
query_lens=query_lens,
|
||||
kv_lens=kv_lens,
|
||||
block_tables=block_tables,
|
||||
scale=scale,
|
||||
sliding_window=sliding_window,
|
||||
soft_cap=soft_cap,
|
||||
)
|
||||
|
||||
output_fp16 = output.to(torch.float32) * output_scale.item()
|
||||
output_fp16 = output_fp16.to(torch.float16)
|
||||
|
||||
atol, rtol = 2e-1, 2e-1
|
||||
(
|
||||
torch.testing.assert_close(output_fp16, ref_output, atol=atol, rtol=rtol),
|
||||
f"{torch.max(torch.abs(output_fp16 - ref_output))}",
|
||||
)
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -21,7 +23,7 @@ QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
|
||||
VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
|
||||
# Avoid combinatorial explosion with full Cartesian product
|
||||
NUM_TOKENS_HIDDEN_SIZES = [
|
||||
*[(1, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5120, 5137]],
|
||||
*[(1, i) for i in [1, 64, 128, *VEC_HIDDEN_SIZES, 5120, 5137]],
|
||||
*[(2048, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5137]],
|
||||
*[(4096, i) for i in [1, 64, 5137]],
|
||||
]
|
||||
@@ -29,6 +31,7 @@ NUM_TOKENS_HIDDEN_SIZES = [
|
||||
ADD_RESIDUAL = [False, True]
|
||||
SCALE_UBS = [True, False]
|
||||
GROUP_SIZES = [None, [1, 64], [1, 128]]
|
||||
TMA_ALIGNMENTS = [0, 4]
|
||||
SEEDS = [0]
|
||||
CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
|
||||
|
||||
@@ -110,12 +113,21 @@ def ops_dynamic_per_token_or_block_quant(
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
if residual is not None:
|
||||
residual = residual.clone()
|
||||
if group_size is not None:
|
||||
out, scales = ops.rms_norm_per_block_quant(
|
||||
x, weight, EPS, quant_dtype, group_size, scale_ub, residual, True
|
||||
x,
|
||||
weight,
|
||||
EPS,
|
||||
quant_dtype,
|
||||
group_size,
|
||||
scale_ub,
|
||||
residual,
|
||||
True,
|
||||
tma_alignment,
|
||||
)
|
||||
scales = scales.contiguous()
|
||||
else:
|
||||
@@ -132,9 +144,10 @@ def ops_impl(
|
||||
residual: torch.Tensor | None,
|
||||
scale_ub: torch.Tensor | None,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
|
||||
return ops_dynamic_per_token_or_block_quant(
|
||||
weight, x, quant_dtype, residual, scale_ub, group_size
|
||||
weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
|
||||
)
|
||||
|
||||
|
||||
@@ -143,7 +156,10 @@ def ops_impl(
|
||||
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
|
||||
@pytest.mark.parametrize("group_size", GROUP_SIZES)
|
||||
@pytest.mark.parametrize(
|
||||
"group_size, tma_alignment",
|
||||
[(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
|
||||
)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
@@ -156,6 +172,7 @@ def test_rms_norm(
|
||||
dtype: torch.dtype,
|
||||
quant_dtype: torch.dtype,
|
||||
group_size: list[int] | None,
|
||||
tma_alignment: int,
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
@@ -173,6 +190,20 @@ def test_rms_norm(
|
||||
# blockwise baseline doesn't support scale_ub
|
||||
return
|
||||
|
||||
if (
|
||||
group_size is None or quant_dtype != current_platform.fp8_dtype()
|
||||
) and tma_alignment != 0:
|
||||
# TMA alignment is only supported for groupwise fp8 kernels
|
||||
return
|
||||
|
||||
if (
|
||||
group_size is not None
|
||||
and tma_alignment != 0
|
||||
and hidden_size // group_size[1] % tma_alignment == 0
|
||||
):
|
||||
# Skip tests where TMA alignment doesn't create extra padding to save time
|
||||
return
|
||||
|
||||
if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
|
||||
# skip
|
||||
return
|
||||
@@ -196,7 +227,7 @@ def test_rms_norm(
|
||||
layer, x, quant_dtype, residual, scale_ub, group_size
|
||||
)
|
||||
ops_out, ops_scales, ops_residual = ops_impl(
|
||||
layer.weight, x, quant_dtype, residual, scale_ub, group_size
|
||||
layer.weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment
|
||||
)
|
||||
|
||||
assert ref_out.dtype == quant_dtype
|
||||
|
||||
@@ -396,6 +396,55 @@ def test_fused_moe(
|
||||
)
|
||||
|
||||
|
||||
def test_fused_moe_int64_overflow(monkeypatch, workspace_init):
|
||||
"""Regression test for int32 overflow in stride*offset products.
|
||||
|
||||
When chunking is disabled and M is large, stride_cm * offs_token can
|
||||
exceed int32 max. Verifies the offs_token int64 cast (fix for #34413)
|
||||
prevents overflow and produces correct results.
|
||||
|
||||
Reproduces the scenario from PR #34279.
|
||||
"""
|
||||
# ~12 GB GPU memory needed for intermediate caches
|
||||
free_mem = torch.cuda.mem_get_info()[0]
|
||||
if free_mem < 12 * 1024**3:
|
||||
pytest.skip("Insufficient GPU memory for overflow test")
|
||||
|
||||
set_random_seed(7)
|
||||
|
||||
m, n, k, e, topk = 100000, 2048, 1024, 8, 6
|
||||
dtype = torch.bfloat16
|
||||
|
||||
# Disable chunking to expose the overflow-prone code path
|
||||
monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "10000000")
|
||||
|
||||
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
|
||||
w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
|
||||
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
|
||||
score = torch.randn((m, e), device="cuda", dtype=dtype)
|
||||
|
||||
# Verify the test exercises the overflow condition:
|
||||
# C has shape (M, topk, N) where N = w1.size(1) = 2*n
|
||||
# stride_cm = C.stride(1) = N, max offs_token = M * topk
|
||||
# Product must exceed int32 max for this test to be meaningful
|
||||
N = w1.size(1)
|
||||
assert N * m * topk > 2**31 - 1, "Test params don't trigger int32 overflow"
|
||||
|
||||
fused_moe_fn = functools.partial(fused_moe, renormalize=False)
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
run_moe_test(
|
||||
torch_moe,
|
||||
fused_moe_fn,
|
||||
a=a,
|
||||
w1=w1,
|
||||
w2=w2,
|
||||
score=score,
|
||||
topk=topk,
|
||||
global_num_experts=e,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m,n,k", FUSED_MOE_MNK_FACTORS_SMALL_M)
|
||||
@pytest.mark.parametrize("e", NUM_EXPERTS_LARGE)
|
||||
@pytest.mark.parametrize("topk", TOP_KS_SMALL)
|
||||
|
||||
@@ -9,6 +9,7 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
|
||||
UnquantizedMoeBackend,
|
||||
select_unquantized_moe_backend,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -65,6 +66,9 @@ def test_select_default_backend_by_platform(
|
||||
"vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
|
||||
return_value=(True, None),
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
|
||||
)
|
||||
def test_select_cuda_flashinfer_trtllm_backend(
|
||||
mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
|
||||
):
|
||||
@@ -101,6 +105,9 @@ def test_select_cuda_flashinfer_trtllm_backend(
|
||||
"vllm.model_executor.layers.fused_moe.oracle.unquantized.is_supported_config_trtllm_bf16",
|
||||
return_value=(False, None),
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(), reason="Only supported on NVIDIA platforms."
|
||||
)
|
||||
def test_select_cuda_flashinfer_cutlass_backend(
|
||||
mock_has_flashinfer, mock_is_supported_trtllm, monkeypatch
|
||||
):
|
||||
|
||||
@@ -169,13 +169,6 @@ VLM_TEST_SETTINGS = {
|
||||
auto_cls=AutoModelForImageTextToText,
|
||||
vllm_output_post_proc=model_utils.qwen2_vllm_to_hf_output,
|
||||
patch_hf_runner=model_utils.qwen3_vl_patch_hf_runner,
|
||||
vllm_runner_kwargs={
|
||||
"attention_config": {
|
||||
"backend": "ROCM_AITER_FA",
|
||||
},
|
||||
}
|
||||
if current_platform.is_rocm()
|
||||
else None,
|
||||
image_size_factors=[(0.25,), (0.25, 0.25, 0.25), (0.25, 0.2, 0.15)],
|
||||
marks=[
|
||||
pytest.mark.core_model,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user