diff --git a/.buildkite/generate_index.py b/.buildkite/generate_index.py
new file mode 100644
index 0000000000000..8350e2705141e
--- /dev/null
+++ b/.buildkite/generate_index.py
@@ -0,0 +1,24 @@
+import argparse
+import os
+
+template = """
+
+
+ Links for vLLM
+ {wheel}
+
+
+"""
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--wheel", help="The wheel path.", required=True)
+args = parser.parse_args()
+
+filename = os.path.basename(args.wheel)
+
+with open("index.html", "w") as f:
+ print(f"Generated index.html for {args.wheel}")
+ # cloudfront requires escaping the '+' character
+ f.write(
+ template.format(wheel=filename,
+ wheel_html_escaped=filename.replace("+", "%2B")))
diff --git a/.buildkite/nightly-benchmarks/benchmark-pipeline.yaml b/.buildkite/nightly-benchmarks/benchmark-pipeline.yaml
index 3db77d5f16022..679abf1814aa5 100644
--- a/.buildkite/nightly-benchmarks/benchmark-pipeline.yaml
+++ b/.buildkite/nightly-benchmarks/benchmark-pipeline.yaml
@@ -1,5 +1,6 @@
steps:
- label: "Wait for container to be ready"
+ key: wait-for-container-image
agents:
queue: A100
plugins:
@@ -10,18 +11,17 @@ steps:
command:
- sh .buildkite/nightly-benchmarks/scripts/wait-for-image.sh
- - wait
-
- label: "A100"
# skip: "use this flag to conditionally skip the benchmark step, useful for PR testing"
agents:
queue: A100
+ depends_on: wait-for-container-image
plugins:
- kubernetes:
podSpec:
priorityClassName: perf-benchmark
containers:
- - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+ - image: public.ecr.aws/q9t5s3a7/vllm-ci-postmerge-repo:$BUILDKITE_COMMIT
command:
- bash .buildkite/nightly-benchmarks/scripts/run-performance-benchmarks.sh
resources:
@@ -49,9 +49,10 @@ steps:
# skip: "use this flag to conditionally skip the benchmark step, useful for PR testing"
agents:
queue: H200
+ depends_on: wait-for-container-image
plugins:
- docker#v5.12.0:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+ image: public.ecr.aws/q9t5s3a7/vllm-ci-postmerge-repo:$BUILDKITE_COMMIT
command:
- bash
- .buildkite/nightly-benchmarks/scripts/run-performance-benchmarks.sh
@@ -65,13 +66,18 @@ steps:
- VLLM_USAGE_SOURCE
- HF_TOKEN
+ #- block: "Run H100 Benchmark"
+ #key: block-h100
+ #depends_on: ~
+
- label: "H100"
# skip: "use this flag to conditionally skip the benchmark step, useful for PR testing"
agents:
queue: H100
+ depends_on: wait-for-container-image
plugins:
- docker#v5.12.0:
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+ image: public.ecr.aws/q9t5s3a7/vllm-ci-postmerge-repo:$BUILDKITE_COMMIT
command:
- bash
- .buildkite/nightly-benchmarks/scripts/run-performance-benchmarks.sh
diff --git a/.buildkite/nightly-benchmarks/scripts/wait-for-image.sh b/.buildkite/nightly-benchmarks/scripts/wait-for-image.sh
index 19f7160e68a4d..aa0f7ade808e0 100644
--- a/.buildkite/nightly-benchmarks/scripts/wait-for-image.sh
+++ b/.buildkite/nightly-benchmarks/scripts/wait-for-image.sh
@@ -1,6 +1,6 @@
#!/bin/sh
-TOKEN=$(curl -s -L "https://public.ecr.aws/token?service=public.ecr.aws&scope=repository:q9t5s3a7/vllm-ci-test-repo:pull" | jq -r .token)
-URL="https://public.ecr.aws/v2/q9t5s3a7/vllm-ci-test-repo/manifests/$BUILDKITE_COMMIT"
+TOKEN=$(curl -s -L "https://public.ecr.aws/token?service=public.ecr.aws&scope=repository:q9t5s3a7/vllm-ci-postmerge-repo:pull" | jq -r .token)
+URL="https://public.ecr.aws/v2/q9t5s3a7/vllm-ci-postmerge-repo/manifests/$BUILDKITE_COMMIT"
TIMEOUT_SECONDS=10
diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml
index f78e360b7afd3..51618a2955fb1 100644
--- a/.buildkite/release-pipeline.yaml
+++ b/.buildkite/release-pipeline.yaml
@@ -1,7 +1,7 @@
steps:
- label: "Build wheel - CUDA 12.1"
agents:
- queue: cpu_queue
+ queue: cpu_queue_postmerge
commands:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.1.0 --tag vllm-ci:build-image --target build --progress plain ."
- "mkdir artifacts"
@@ -18,7 +18,7 @@ steps:
- label: "Build wheel - CUDA 11.8"
# depends_on: block-build-cu118-wheel
agents:
- queue: cpu_queue
+ queue: cpu_queue_postmerge
commands:
- "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=11.8.0 --tag vllm-ci:build-image --target build --progress plain ."
- "mkdir artifacts"
@@ -26,3 +26,47 @@ steps:
- "bash .buildkite/upload-wheels.sh"
env:
DOCKER_BUILDKIT: "1"
+
+ - block: "Build release image"
+ depends_on: ~
+ key: block-release-image-build
+
+ - label: "Build release image"
+ depends_on: block-release-image-build
+ agents:
+ queue: cpu_queue_postmerge
+ commands:
+ - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+ - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.1.0 --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT --target vllm-openai --progress plain ."
+ - "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
+
+ - label: "Build and publish TPU release image"
+ depends_on: ~
+ if: build.env("NIGHTLY") == "1"
+ agents:
+ queue: tpu_queue_postmerge
+ commands:
+ - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --tag vllm/vllm-tpu:nightly --tag vllm/vllm-tpu:$BUILDKITE_COMMIT --progress plain -f Dockerfile.tpu ."
+ - "docker push vllm/vllm-tpu:nightly"
+ - "docker push vllm/vllm-tpu:$BUILDKITE_COMMIT"
+ plugins:
+ - docker-login#v3.0.0:
+ username: vllm
+ password-env: DOCKERHUB_TOKEN
+ env:
+ DOCKER_BUILDKIT: "1"
+
+ - block: "Build CPU release image"
+ key: block-cpu-release-image-build
+ depends_on: ~
+
+ - label: "Build and publish CPU release image"
+ depends_on: block-cpu-release-image-build
+ agents:
+ queue: cpu_queue_postmerge
+ commands:
+ - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+ - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg GIT_REPO_CHECK=1 --tag public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$RELEASE_VERSION --progress plain -f Dockerfile.cpu ."
+ - "docker push public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:$RELEASE_VERSION"
+ env:
+ DOCKER_BUILDKIT: "1"
diff --git a/.buildkite/run-cpu-test.sh b/.buildkite/run-cpu-test.sh
index 4f1729d46dae2..9925db7bea593 100644
--- a/.buildkite/run-cpu-test.sh
+++ b/.buildkite/run-cpu-test.sh
@@ -9,36 +9,33 @@ CORE_RANGE=${CORE_RANGE:-48-95}
NUMA_NODE=${NUMA_NODE:-1}
# Try building the docker image
-numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build -t cpu-test -f Dockerfile.cpu .
-numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --build-arg VLLM_CPU_DISABLE_AVX512="true" -t cpu-test-avx2 -f Dockerfile.cpu .
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build -t cpu-test-"$BUILDKITE_BUILD_NUMBER" -f Dockerfile.cpu .
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --build-arg VLLM_CPU_DISABLE_AVX512="true" -t cpu-test-"$BUILDKITE_BUILD_NUMBER"-avx2 -f Dockerfile.cpu .
# Setup cleanup
-remove_docker_container() { docker rm -f cpu-test-"$NUMA_NODE" cpu-test-avx2-"$NUMA_NODE" || true; }
+remove_docker_container() { set -e; docker rm -f cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" cpu-test-"$BUILDKITE_BUILD_NUMBER"-avx2-"$NUMA_NODE" || true; }
trap remove_docker_container EXIT
remove_docker_container
# Run the image, setting --shm-size=4g for tensor parallel.
docker run -itd --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --cpuset-cpus="$CORE_RANGE" \
- --cpuset-mems="$NUMA_NODE" --privileged=true --network host -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --shm-size=4g --name cpu-test-"$NUMA_NODE" cpu-test
+ --cpuset-mems="$NUMA_NODE" --privileged=true --network host -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --shm-size=4g --name cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" cpu-test-"$BUILDKITE_BUILD_NUMBER"
docker run -itd --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --cpuset-cpus="$CORE_RANGE" \
- --cpuset-mems="$NUMA_NODE" --privileged=true --network host -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --shm-size=4g --name cpu-test-avx2-"$NUMA_NODE" cpu-test-avx2
+ --cpuset-mems="$NUMA_NODE" --privileged=true --network host -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --shm-size=4g --name cpu-test-"$BUILDKITE_BUILD_NUMBER"-avx2-"$NUMA_NODE" cpu-test-"$BUILDKITE_BUILD_NUMBER"-avx2
function cpu_tests() {
set -e
export NUMA_NODE=$2
# offline inference
- docker exec cpu-test-avx2-"$NUMA_NODE" bash -c "
+ docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-avx2-"$NUMA_NODE" bash -c "
set -e
- python3 examples/offline_inference.py"
+ python3 examples/offline_inference/basic.py"
# Run basic model test
- docker exec cpu-test-"$NUMA_NODE" bash -c "
+ docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
set -e
- pip install pytest pytest-asyncio \
- decord einops librosa peft Pillow sentence-transformers soundfile \
- transformers_stream_generator matplotlib datamodel_code_generator
- pip install torchvision --index-url https://download.pytorch.org/whl/cpu
+ pip install -r vllm/requirements-test.txt
pytest -v -s tests/models/decoder_only/language -m cpu_model
pytest -v -s tests/models/embedding/language -m cpu_model
pytest -v -s tests/models/encoder_decoder/language -m cpu_model
@@ -46,26 +43,26 @@ function cpu_tests() {
pytest -v -s tests/models/decoder_only/vision_language -m cpu_model"
# Run compressed-tensor test
- docker exec cpu-test-"$NUMA_NODE" bash -c "
+ docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
set -e
pytest -s -v \
tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_static_setup \
tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_dynamic_per_token"
# Run AWQ test
- docker exec cpu-test-"$NUMA_NODE" bash -c "
+ docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
set -e
pytest -s -v \
tests/quantization/test_ipex_quant.py"
# Run chunked-prefill and prefix-cache test
- docker exec cpu-test-"$NUMA_NODE" bash -c "
+ docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
set -e
pytest -s -v -k cpu_model \
tests/basic_correctness/test_chunked_prefill.py"
- # online inference
- docker exec cpu-test-"$NUMA_NODE" bash -c "
+ # online serving
+ docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
set -e
export VLLM_CPU_KVCACHE_SPACE=10
export VLLM_CPU_OMP_THREADS_BIND=$1
@@ -78,6 +75,12 @@ function cpu_tests() {
--num-prompts 20 \
--endpoint /v1/completions \
--tokenizer facebook/opt-125m"
+
+ # Run multi-lora tests
+ docker exec cpu-test-"$BUILDKITE_BUILD_NUMBER"-"$NUMA_NODE" bash -c "
+ set -e
+ pytest -s -v \
+ tests/lora/test_qwen2vl.py"
}
# All of CPU tests are expected to be finished less than 25 mins.
diff --git a/.buildkite/run-gh200-test.sh b/.buildkite/run-gh200-test.sh
new file mode 100644
index 0000000000000..3e4e409466b8a
--- /dev/null
+++ b/.buildkite/run-gh200-test.sh
@@ -0,0 +1,28 @@
+#!/bin/bash
+
+# This script build the GH200 docker image and run the offline inference inside the container.
+# It serves a sanity check for compilation and basic model usage.
+set -ex
+
+# Skip the new torch installation during build since we are using the specified version for arm64 in the Dockerfile
+python3 use_existing_torch.py
+
+# Try building the docker image
+DOCKER_BUILDKIT=1 docker build . \
+ --target vllm-openai \
+ --platform "linux/arm64" \
+ -t gh200-test \
+ --build-arg max_jobs=66 \
+ --build-arg nvcc_threads=2 \
+ --build-arg torch_cuda_arch_list="9.0+PTX" \
+ --build-arg vllm_fa_cmake_gpu_arches="90-real"
+
+# Setup cleanup
+remove_docker_container() { docker rm -f gh200-test || true; }
+trap remove_docker_container EXIT
+remove_docker_container
+
+# Run the image and test offline inference
+docker run --name gh200-test --gpus=all --entrypoint="" gh200-test bash -c '
+ python3 examples/offline_inference/basic.py
+'
diff --git a/.buildkite/run-hpu-test.sh b/.buildkite/run-hpu-test.sh
index fa4f74fca7a11..8f3b08212fd6a 100644
--- a/.buildkite/run-hpu-test.sh
+++ b/.buildkite/run-hpu-test.sh
@@ -13,4 +13,4 @@ trap remove_docker_container EXIT
remove_docker_container
# Run the image and launch offline inference
-docker run --runtime=habana --name=hpu-test --network=host -e HABANA_VISIBLE_DEVICES=all -e VLLM_SKIP_WARMUP=true --entrypoint="" hpu-test-env python3 examples/offline_inference.py
\ No newline at end of file
+docker run --runtime=habana --name=hpu-test --network=host -e HABANA_VISIBLE_DEVICES=all -e VLLM_SKIP_WARMUP=true --entrypoint="" hpu-test-env python3 examples/offline_inference/basic.py
\ No newline at end of file
diff --git a/.buildkite/run-neuron-test.sh b/.buildkite/run-neuron-test.sh
index 9259391aaed49..189714ebb6d75 100644
--- a/.buildkite/run-neuron-test.sh
+++ b/.buildkite/run-neuron-test.sh
@@ -3,6 +3,18 @@
# This script build the Neuron docker image and run the API server inside the container.
# It serves a sanity check for compilation and basic model usage.
set -e
+set -v
+
+image_name="neuron/vllm-ci"
+container_name="neuron_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
+
+HF_CACHE="$(realpath ~)/huggingface"
+mkdir -p "${HF_CACHE}"
+HF_MOUNT="/root/.cache/huggingface"
+
+NEURON_COMPILE_CACHE_URL="$(realpath ~)/neuron_compile_cache"
+mkdir -p "${NEURON_COMPILE_CACHE_URL}"
+NEURON_COMPILE_CACHE_MOUNT="/root/.cache/neuron_compile_cache"
# Try building the docker image
aws ecr get-login-password --region us-west-2 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-west-2.amazonaws.com
@@ -13,41 +25,30 @@ if [ -f /tmp/neuron-docker-build-timestamp ]; then
last_build=$(cat /tmp/neuron-docker-build-timestamp)
current_time=$(date +%s)
if [ $((current_time - last_build)) -gt 86400 ]; then
+ docker image prune -f
docker system prune -f
+ rm -rf "${HF_MOUNT:?}/*"
+ rm -rf "${NEURON_COMPILE_CACHE_MOUNT:?}/*"
echo "$current_time" > /tmp/neuron-docker-build-timestamp
fi
else
date "+%s" > /tmp/neuron-docker-build-timestamp
fi
-docker build -t neuron -f Dockerfile.neuron .
+docker build -t "${image_name}" -f Dockerfile.neuron .
# Setup cleanup
-remove_docker_container() { docker rm -f neuron || true; }
+remove_docker_container() {
+ docker image rm -f "${image_name}" || true;
+}
trap remove_docker_container EXIT
-remove_docker_container
# Run the image
-docker run --device=/dev/neuron0 --device=/dev/neuron1 --network host --name neuron neuron python3 -m vllm.entrypoints.api_server \
- --model TinyLlama/TinyLlama-1.1B-Chat-v1.0 --max-num-seqs 8 --max-model-len 128 --block-size 128 --device neuron --tensor-parallel-size 2 &
-
-# Wait for the server to start
-wait_for_server_to_start() {
- timeout=300
- counter=0
-
- while [ "$(curl -s -o /dev/null -w '%{http_code}' localhost:8000/health)" != "200" ]; do
- sleep 1
- counter=$((counter + 1))
- if [ $counter -ge $timeout ]; then
- echo "Timeout after $timeout seconds"
- break
- fi
- done
-}
-wait_for_server_to_start
-
-# Test a simple prompt
-curl -X POST -H "Content-Type: application/json" \
- localhost:8000/generate \
- -d '{"prompt": "San Francisco is a"}'
+docker run --rm -it --device=/dev/neuron0 --device=/dev/neuron1 --network host \
+ -v "${HF_CACHE}:${HF_MOUNT}" \
+ -e "HF_HOME=${HF_MOUNT}" \
+ -v "${NEURON_COMPILE_CACHE_URL}:${NEURON_COMPILE_CACHE_MOUNT}" \
+ -e "NEURON_COMPILE_CACHE_URL=${NEURON_COMPILE_CACHE_MOUNT}" \
+ --name "${container_name}" \
+ ${image_name} \
+ /bin/bash -c "python3 /workspace/vllm/examples/offline_inference/neuron.py"
diff --git a/.buildkite/run-openvino-test.sh b/.buildkite/run-openvino-test.sh
index 6b12f424fd828..6159b21ff8206 100755
--- a/.buildkite/run-openvino-test.sh
+++ b/.buildkite/run-openvino-test.sh
@@ -13,4 +13,4 @@ trap remove_docker_container EXIT
remove_docker_container
# Run the image and launch offline inference
-docker run --network host --env VLLM_OPENVINO_KVCACHE_SPACE=1 --name openvino-test openvino-test python3 /workspace/examples/offline_inference.py
+docker run --network host --env VLLM_OPENVINO_KVCACHE_SPACE=1 --name openvino-test openvino-test python3 /workspace/examples/offline_inference/basic.py
diff --git a/.buildkite/run-tpu-test.sh b/.buildkite/run-tpu-test.sh
index 770dad6ffa3a1..650af0fac4c61 100644
--- a/.buildkite/run-tpu-test.sh
+++ b/.buildkite/run-tpu-test.sh
@@ -14,4 +14,13 @@ remove_docker_container
# For HF_TOKEN.
source /etc/environment
# Run a simple end-to-end example.
-docker run --privileged --net host --shm-size=16G -it -e "HF_TOKEN=$HF_TOKEN" --name tpu-test vllm-tpu /bin/bash -c "python3 -m pip install git+https://github.com/thuml/depyf.git && python3 -m pip install pytest && python3 -m pip install lm_eval[api]==0.4.4 && pytest -v -s /workspace/vllm/tests/entrypoints/openai/test_accuracy.py && pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py && python3 /workspace/vllm/tests/tpu/test_compilation.py && python3 /workspace/vllm/examples/offline_inference_tpu.py"
+docker run --privileged --net host --shm-size=16G -it \
+ -e "HF_TOKEN=$HF_TOKEN" --name tpu-test \
+ vllm-tpu /bin/bash -c "python3 -m pip install git+https://github.com/thuml/depyf.git \
+ && python3 -m pip install pytest \
+ && python3 -m pip install lm_eval[api]==0.4.4 \
+ && pytest -v -s /workspace/vllm/tests/entrypoints/openai/test_accuracy.py \
+ && pytest -v -s /workspace/vllm/tests/tpu/test_custom_dispatcher.py \
+ && python3 /workspace/vllm/tests/tpu/test_compilation.py \
+ && python3 /workspace/vllm/tests/tpu/test_quantization_accuracy.py \
+ && python3 /workspace/vllm/examples/offline_inference/tpu.py"
diff --git a/.buildkite/run-xpu-test.sh b/.buildkite/run-xpu-test.sh
index faeac8e2ded36..4d344e58db8ac 100644
--- a/.buildkite/run-xpu-test.sh
+++ b/.buildkite/run-xpu-test.sh
@@ -12,5 +12,8 @@ remove_docker_container() { docker rm -f xpu-test || true; }
trap remove_docker_container EXIT
remove_docker_container
-# Run the image and launch offline inference
-docker run --network host --name xpu-test --device /dev/dri -v /dev/dri/by-path:/dev/dri/by-path --entrypoint="" xpu-test python3 examples/offline_inference.py
+# Run the image and test offline inference/tensor parallel
+docker run --name xpu-test --device /dev/dri -v /dev/dri/by-path:/dev/dri/by-path --entrypoint="" xpu-test sh -c '
+ python3 examples/offline_inference/basic.py
+ python3 examples/offline_inference/cli.py -tp 2
+'
diff --git a/.buildkite/test-pipeline.yaml b/.buildkite/test-pipeline.yaml
index f5591f1098534..74b287c7adbfa 100644
--- a/.buildkite/test-pipeline.yaml
+++ b/.buildkite/test-pipeline.yaml
@@ -38,7 +38,7 @@ steps:
- pip install -r requirements-docs.txt
- SPHINXOPTS=\"-W\" make html
# Check API reference (if it fails, you may have missing mock imports)
- - grep \"sig sig-object py\" build/html/dev/sampling_params.html
+ - grep \"sig sig-object py\" build/html/api/inference_params.html
- label: Async Engine, Inputs, Utils, Worker Test # 24min
fast_check: true
@@ -50,9 +50,10 @@ steps:
- tests/multimodal
- tests/test_utils
- tests/worker
- - tests/test_lazy_torch_compile.py
+ - tests/standalone_tests/lazy_torch_compile.py
commands:
- - python3 test_lazy_torch_compile.py
+ - pip install git+https://github.com/Isotr0py/DeepSeek-VL2.git # Used by multimoda processing test
+ - python3 standalone_tests/lazy_torch_compile.py
- pytest -v -s mq_llm_engine # MQLLMEngine
- pytest -v -s async_engine # AsyncLLMEngine
- NUM_SCHEDULER_STEPS=4 pytest -v -s async_engine/test_async_llm_engine.py
@@ -61,6 +62,13 @@ steps:
- pytest -v -s test_utils.py # Utils
- pytest -v -s worker # Worker
+- label: Python-only Installation Test
+ source_file_dependencies:
+ - tests/standalone_tests/python_only_compile.sh
+ - setup.py
+ commands:
+ - bash standalone_tests/python_only_compile.sh
+
- label: Basic Correctness Test # 30min
#mirror_hardwares: [amd]
fast_check: true
@@ -99,14 +107,12 @@ steps:
source_file_dependencies:
- vllm/
commands:
- - pip install -e ./plugins/vllm_add_dummy_model
- pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_lazy_outlines.py --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_generate_multiple_loras.py --ignore=entrypoints/llm/test_guided_generate.py
- pytest -v -s entrypoints/llm/test_lazy_outlines.py # it needs a clean process
- pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process
- pytest -v -s entrypoints/llm/test_generate_multiple_loras.py # it needs a clean process
- pytest -v -s entrypoints/llm/test_guided_generate.py # it needs a clean process
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_oot_registration.py
- - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
- pytest -v -s entrypoints/test_chat_utils.py
- pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
@@ -174,27 +180,30 @@ steps:
commands:
- VLLM_USE_V1=1 pytest -v -s v1
-- label: Examples Test # 15min
+- label: Examples Test # 25min
working_dir: "/vllm-workspace/examples"
#mirror_hardwares: [amd]
source_file_dependencies:
- vllm/entrypoints
- examples/
commands:
- - pip install awscli tensorizer # for llava example and tensorizer test
- - python3 offline_inference.py
- - python3 cpu_offload.py
- - python3 offline_inference_chat.py
- - python3 offline_inference_with_prefix.py
- - python3 llm_engine_example.py
- - python3 offline_inference_vision_language.py
- - python3 offline_inference_vision_language_multi_image.py
- - python3 tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
- - python3 offline_inference_encoder_decoder.py
- - python3 offline_profile.py --model facebook/opt-125m
+ - pip install tensorizer # for tensorizer test
+ - python3 offline_inference/basic.py
+ - python3 offline_inference/cpu_offload.py
+ - python3 offline_inference/chat.py
+ - python3 offline_inference/prefix_caching.py
+ - python3 offline_inference/llm_engine_example.py
+ - python3 offline_inference/vision_language.py
+ - python3 offline_inference/vision_language_multi_image.py
+ - python3 other/tensorize_vllm_model.py --model facebook/opt-125m serialize --serialized-directory /tmp/ --suffix v1 && python3 other/tensorize_vllm_model.py --model facebook/opt-125m deserialize --path-to-tensors /tmp/vllm/facebook/opt-125m/v1/model.tensors
+ - python3 offline_inference/encoder_decoder.py
+ - python3 offline_inference/classification.py
+ - python3 offline_inference/embedding.py
+ - python3 offline_inference/scoring.py
+ - python3 offline_inference/profiling.py --model facebook/opt-125m run_num_steps --num-steps 2
- label: Prefix Caching Test # 9min
- #mirror_hardwares: [amd]
+ mirror_hardwares: [amd]
source_file_dependencies:
- vllm/
- tests/prefix_caching
@@ -206,6 +215,7 @@ steps:
- vllm/model_executor/layers
- vllm/sampling_metadata.py
- tests/samplers
+ - tests/conftest.py
commands:
- pytest -v -s samplers
- VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
@@ -214,23 +224,29 @@ steps:
mirror_hardwares: [amd]
source_file_dependencies:
- vllm/model_executor/layers
+ - vllm/model_executor/guided_decoding
- tests/test_logits_processor
- command: pytest -v -s test_logits_processor.py
+ - tests/model_executor/test_guided_processors
+ commands:
+ - pytest -v -s test_logits_processor.py
+ - pytest -v -s model_executor/test_guided_processors.py
-- label: Speculative decoding tests # 30min
+- label: Speculative decoding tests # 40min
source_file_dependencies:
- vllm/spec_decode
- tests/spec_decode
+ - vllm/model_executor/models/eagle.py
commands:
- pytest -v -s spec_decode/e2e/test_multistep_correctness.py
- VLLM_ATTENTION_BACKEND=FLASH_ATTN pytest -v -s spec_decode --ignore=spec_decode/e2e/test_multistep_correctness.py
+ - pytest -v -s spec_decode/e2e/test_eagle_correctness.py
- label: LoRA Test %N # 15min each
mirror_hardwares: [amd]
source_file_dependencies:
- vllm/lora
- tests/lora
- command: pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore lora/test_long_context.py lora/test_chatglm3_tp.py lora/test_llama_tp.py
+ command: pytest -v -s lora --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --ignore=lora/test_long_context.py --ignore=lora/test_chatglm3_tp.py --ignore=lora/test_llama_tp.py --ignore=lora/test_minicpmv_tp.py
parallelism: 4
- label: "PyTorch Fullgraph Smoke Test" # 9min
@@ -314,17 +330,15 @@ steps:
##### models test #####
-- label: Basic Models Test # 30min
+- label: Basic Models Test # 24min
source_file_dependencies:
- vllm/
- tests/models
commands:
- - pip install -e ./plugins/vllm_add_dummy_model
- - pytest -v -s models/test_oot_registration.py # it needs a clean process
- pytest -v -s models/test_registry.py
- pytest -v -s models/test_initialization.py
-- label: Language Models Test (Standard) # 42min
+- label: Language Models Test (Standard) # 32min
#mirror_hardwares: [amd]
source_file_dependencies:
- vllm/
@@ -335,7 +349,7 @@ steps:
- pytest -v -s models/decoder_only/language -m 'core_model or quant_model'
- pytest -v -s models/embedding/language -m core_model
-- label: Language Models Test (Extended) # 50min
+- label: Language Models Test (Extended) # 1h10min
optional: true
source_file_dependencies:
- vllm/
@@ -346,22 +360,26 @@ steps:
- pytest -v -s models/decoder_only/language -m 'not core_model and not quant_model'
- pytest -v -s models/embedding/language -m 'not core_model'
-- label: Multi-Modal Models Test (Standard) # 26min
+- label: Multi-Modal Models Test (Standard) # 40min
#mirror_hardwares: [amd]
source_file_dependencies:
- vllm/
- tests/models/decoder_only/audio_language
- tests/models/decoder_only/vision_language
- tests/models/embedding/vision_language
+ - tests/models/encoder_decoder/audio_language
- tests/models/encoder_decoder/vision_language
commands:
+ - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
+ - pytest -v -s models/multimodal
- pytest -v -s models/decoder_only/audio_language -m 'core_model or quant_model'
- pytest -v -s --ignore models/decoder_only/vision_language/test_phi3v.py models/decoder_only/vision_language -m 'core_model or quant_model'
- pytest -v -s models/embedding/vision_language -m core_model
+ - pytest -v -s models/encoder_decoder/audio_language -m core_model
- pytest -v -s models/encoder_decoder/language -m core_model
- pytest -v -s models/encoder_decoder/vision_language -m core_model
-- label: Multi-Modal Models Test (Extended) # 1h15m
+- label: Multi-Modal Models Test (Extended) 1 # 48m
optional: true
source_file_dependencies:
- vllm/
@@ -370,15 +388,26 @@ steps:
- tests/models/embedding/vision_language
- tests/models/encoder_decoder/vision_language
commands:
+ - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
- pytest -v -s models/decoder_only/audio_language -m 'not core_model and not quant_model'
+ - pytest -v -s models/decoder_only/vision_language/test_models.py -m 'split(group=0) and not core_model and not quant_model'
# HACK - run phi3v tests separately to sidestep this transformers bug
# https://github.com/huggingface/transformers/issues/34307
- pytest -v -s models/decoder_only/vision_language/test_phi3v.py
- - pytest -v -s --ignore models/decoder_only/vision_language/test_phi3v.py models/decoder_only/vision_language -m 'not core_model and not quant_model'
+ - pytest -v -s --ignore models/decoder_only/vision_language/test_models.py --ignore models/decoder_only/vision_language/test_phi3v.py models/decoder_only/vision_language -m 'not core_model and not quant_model'
- pytest -v -s models/embedding/vision_language -m 'not core_model'
- pytest -v -s models/encoder_decoder/language -m 'not core_model'
- pytest -v -s models/encoder_decoder/vision_language -m 'not core_model'
+- label: Multi-Modal Models Test (Extended) 2 # 38m
+ optional: true
+ source_file_dependencies:
+ - vllm/
+ - tests/models/decoder_only/vision_language
+ commands:
+ - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
+ - pytest -v -s models/decoder_only/vision_language/test_models.py -m 'split(group=1) and not core_model and not quant_model'
+
# This test is used only in PR development phase to test individual models and should never run on main
- label: Custom Models Test
optional: true
@@ -413,11 +442,11 @@ steps:
- tests/distributed/
commands:
- # the following commands are for the first node, with ip 192.168.10.10 (ray environment already set up)
- - VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep -q 'Same node test passed'
+ - VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
- VLLM_MULTI_NODE=1 pytest -v -s distributed/test_multi_node_assignment.py
- VLLM_MULTI_NODE=1 pytest -v -s distributed/test_pipeline_parallel.py
- # the following commands are for the second node, with ip 192.168.10.11 (ray environment already set up)
- - VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep -q 'Same node test passed'
+ - VLLM_TEST_SAME_HOST=0 torchrun --nnodes 2 --nproc-per-node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.10.10 distributed/test_same_node.py | grep 'Same node test passed'
- label: Distributed Tests (2 GPUs) # 40min
#mirror_hardwares: [amd]
@@ -436,18 +465,35 @@ steps:
commands:
- pytest -v -s ./compile/test_basic_correctness.py
- pytest -v -s ./compile/test_wrapper.py
- - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep -q 'Same node test passed'
- - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m distributed_2_gpus
+ - VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
+ - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
# Avoid importing model tests that cause CUDA reinitialization error
- - pytest models/encoder_decoder/language/test_bart.py -v -s -m distributed_2_gpus
- - pytest models/encoder_decoder/vision_language/test_broadcast.py -v -s -m distributed_2_gpus
- - pytest models/decoder_only/vision_language/test_models.py -v -s -m distributed_2_gpus
+ - pytest models/encoder_decoder/language/test_bart.py -v -s -m 'distributed(num_gpus=2)'
+ - pytest models/encoder_decoder/vision_language/test_broadcast.py -v -s -m 'distributed(num_gpus=2)'
+ - pytest models/decoder_only/vision_language/test_models.py -v -s -m 'distributed(num_gpus=2)'
- pytest -v -s spec_decode/e2e/test_integration_dist_tp2.py
- - pip install -e ./plugins/vllm_add_dummy_model
- - pytest -v -s distributed/test_distributed_oot.py
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s test_sharded_state_loader.py
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s kv_transfer/disagg_test.py
+- label: Plugin Tests (2 GPUs) # 40min
+ working_dir: "/vllm-workspace/tests"
+ num_gpus: 2
+ fast_check: true
+ source_file_dependencies:
+ - vllm/plugins/
+ - tests/plugins/
+ commands:
+ # begin platform plugin tests, all the code in-between runs on dummy platform
+ - pip install -e ./plugins/vllm_add_dummy_platform
+ - pytest -v -s plugins_tests/test_platform_plugins.py
+ - pip uninstall vllm_add_dummy_platform -y
+ # end platform plugin tests
+ # other tests continue here:
+ - pip install -e ./plugins/vllm_add_dummy_model
+ - pytest -v -s distributed/test_distributed_oot.py
+ - pytest -v -s entrypoints/openai/test_oot_registration.py # it needs a clean process
+ - pytest -v -s models/test_oot_registration.py # it needs a clean process
+
- label: Multi-step Tests (4 GPUs) # 36min
working_dir: "/vllm-workspace/tests"
num_gpus: 4
@@ -481,7 +527,6 @@ steps:
- label: LoRA TP Test (Distributed)
num_gpus: 4
- soft_fail: true
source_file_dependencies:
- vllm/lora
- tests/lora
@@ -495,6 +540,7 @@ steps:
# requires multi-GPU testing for validation.
- pytest -v -s -x lora/test_chatglm3_tp.py
- pytest -v -s -x lora/test_llama_tp.py
+ - pytest -v -s -x lora/test_minicpmv_tp.py
- label: Weight Loading Multiple GPU Test # 33min
@@ -532,7 +578,7 @@ steps:
# see https://github.com/vllm-project/vllm/pull/5689 for details
- pytest -v -s distributed/test_custom_all_reduce.py
- torchrun --nproc_per_node=2 distributed/test_ca_buffer_sharing.py
- - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m distributed_2_gpus
+ - TARGET_TEST_SUITE=A100 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
- pytest -v -s -x lora/test_mixtral.py
- label: LM Eval Large Models # optional
diff --git a/.buildkite/upload-wheels.sh b/.buildkite/upload-wheels.sh
index 7345dd4e66b29..3c756659a715a 100644
--- a/.buildkite/upload-wheels.sh
+++ b/.buildkite/upload-wheels.sh
@@ -23,6 +23,8 @@ wheel="$new_wheel"
version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
echo "Version: $version"
+normal_wheel="$wheel" # Save the original wheel filename
+
# If the version contains "dev", rename it to v1.0.0.dev for consistency
if [[ $version == *dev* ]]; then
suffix="${version##*.}"
@@ -32,12 +34,38 @@ if [[ $version == *dev* ]]; then
new_version="1.0.0.dev"
fi
new_wheel="${wheel/$version/$new_version}"
- mv -- "$wheel" "$new_wheel"
+ # use cp to keep both files in the artifacts directory
+ cp -- "$wheel" "$new_wheel"
wheel="$new_wheel"
version="$new_version"
fi
# Upload the wheel to S3
+python3 .buildkite/generate_index.py --wheel "$normal_wheel"
+
+# generate index for this commit
aws s3 cp "$wheel" "s3://vllm-wheels/$BUILDKITE_COMMIT/"
+aws s3 cp "$normal_wheel" "s3://vllm-wheels/$BUILDKITE_COMMIT/"
+
+if [[ $normal_wheel == *"cu118"* ]]; then
+ # if $normal_wheel matches cu118, do not upload the index.html
+ echo "Skipping index files for cu118 wheels"
+else
+ # only upload index.html for cu12 wheels (default wheels)
+ aws s3 cp index.html "s3://vllm-wheels/$BUILDKITE_COMMIT/vllm/index.html"
+ aws s3 cp "s3://vllm-wheels/nightly/index.html" "s3://vllm-wheels/$BUILDKITE_COMMIT/index.html"
+fi
+
+# generate index for nightly
aws s3 cp "$wheel" "s3://vllm-wheels/nightly/"
+aws s3 cp "$normal_wheel" "s3://vllm-wheels/nightly/"
+
+if [[ $normal_wheel == *"cu118"* ]]; then
+ # if $normal_wheel matches cu118, do not upload the index.html
+ echo "Skipping index files for cu118 wheels"
+else
+ # only upload index.html for cu12 wheels (default wheels)
+ aws s3 cp index.html "s3://vllm-wheels/nightly/vllm/index.html"
+fi
+
aws s3 cp "$wheel" "s3://vllm-wheels/$version/"
\ No newline at end of file
diff --git a/.github/ISSUE_TEMPLATE/400-bug report.yml b/.github/ISSUE_TEMPLATE/400-bug-report.yml
similarity index 100%
rename from .github/ISSUE_TEMPLATE/400-bug report.yml
rename to .github/ISSUE_TEMPLATE/400-bug-report.yml
diff --git a/.github/ISSUE_TEMPLATE/500-feature request.yml b/.github/ISSUE_TEMPLATE/500-feature-request.yml
similarity index 100%
rename from .github/ISSUE_TEMPLATE/500-feature request.yml
rename to .github/ISSUE_TEMPLATE/500-feature-request.yml
diff --git a/.github/ISSUE_TEMPLATE/600-new model.yml b/.github/ISSUE_TEMPLATE/600-new-model.yml
similarity index 94%
rename from .github/ISSUE_TEMPLATE/600-new model.yml
rename to .github/ISSUE_TEMPLATE/600-new-model.yml
index 794617a0cfdf6..713e76c1a5cec 100644
--- a/.github/ISSUE_TEMPLATE/600-new model.yml
+++ b/.github/ISSUE_TEMPLATE/600-new-model.yml
@@ -9,7 +9,7 @@ body:
value: >
#### Before submitting an issue, please make sure the issue hasn't been already addressed by searching through [the existing and past issues](https://github.com/vllm-project/vllm/issues?q=is%3Aissue+sort%3Acreated-desc+).
- #### We also highly recommend you read https://docs.vllm.ai/en/latest/models/adding_model.html first to understand how to add a new model.
+ #### We also highly recommend you read https://docs.vllm.ai/en/latest/contributing/model/adding_model.html first to understand how to add a new model.
- type: textarea
attributes:
label: The model to consider.
diff --git a/.github/ISSUE_TEMPLATE/700-performance discussion.yml b/.github/ISSUE_TEMPLATE/700-performance-discussion.yml
similarity index 100%
rename from .github/ISSUE_TEMPLATE/700-performance discussion.yml
rename to .github/ISSUE_TEMPLATE/700-performance-discussion.yml
diff --git a/.github/ISSUE_TEMPLATE/800-misc discussion.yml b/.github/ISSUE_TEMPLATE/800-misc-discussion.yml
similarity index 100%
rename from .github/ISSUE_TEMPLATE/800-misc discussion.yml
rename to .github/ISSUE_TEMPLATE/800-misc-discussion.yml
diff --git a/.github/workflows/sphinx-lint.yml b/.github/workflows/doc-lint.yml
similarity index 94%
rename from .github/workflows/sphinx-lint.yml
rename to .github/workflows/doc-lint.yml
index 255c897fd1891..e66e622ab2c9e 100644
--- a/.github/workflows/sphinx-lint.yml
+++ b/.github/workflows/doc-lint.yml
@@ -15,7 +15,7 @@ on:
- "docs/**"
jobs:
- sphinx-lint:
+ doc-lint:
runs-on: ubuntu-latest
strategy:
matrix:
@@ -31,4 +31,4 @@ jobs:
python -m pip install --upgrade pip
pip install -r requirements-lint.txt
- name: Linting docs
- run: tools/sphinx-lint.sh
+ run: tools/doc-lint.sh
diff --git a/.gitignore b/.gitignore
index ceef6a5fba456..89dab8f13bab1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -79,8 +79,7 @@ instance/
# Sphinx documentation
docs/_build/
-docs/source/getting_started/examples/*.rst
-!**/*.template.rst
+docs/source/getting_started/examples/
# PyBuilder
.pybuilder/
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 38dcf9a591bff..84194a2ff5116 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -230,6 +230,7 @@ set(VLLM_EXT_SRC
"csrc/quantization/gptq/q_gemm.cu"
"csrc/quantization/compressed_tensors/int8_quant_kernels.cu"
"csrc/quantization/fp8/common.cu"
+ "csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu"
"csrc/quantization/gguf/gguf_kernel.cu"
"csrc/cuda_utils_kernels.cu"
"csrc/prepare_inputs/advance_step.cu"
@@ -239,7 +240,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")
# Set CUTLASS_REVISION manually -- its revision detection doesn't work in this case.
- set(CUTLASS_REVISION "v3.5.1" CACHE STRING "CUTLASS revision to use")
+ set(CUTLASS_REVISION "v3.6.0" CACHE STRING "CUTLASS revision to use")
# Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided
if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})
@@ -256,7 +257,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
FetchContent_Declare(
cutlass
GIT_REPOSITORY https://github.com/nvidia/cutlass.git
- GIT_TAG v3.5.1
+ GIT_TAG v3.6.0
GIT_PROGRESS TRUE
# Speed up CUTLASS download by retrieving only the specified GIT_TAG instead of the history.
@@ -274,7 +275,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
"csrc/quantization/awq/gemm_kernels.cu"
"csrc/custom_all_reduce.cu"
"csrc/permute_cols.cu"
- "csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu")
+ "csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu"
+ "csrc/sparse/cutlass/sparse_scaled_mm_entry.cu"
+ "csrc/sparse/cutlass/sparse_compressor_entry.cu"
+ "csrc/cutlass_extensions/common.cpp")
set_gencode_flags_for_srcs(
SRCS "${VLLM_EXT_SRC}"
@@ -303,7 +307,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
" in CUDA target architectures")
endif()
- #
# The cutlass_scaled_mm kernels for Hopper (c3x, i.e. CUTLASS 3.x) require
# CUDA 12.0 or later (and only work on Hopper, 9.0/9.0a for now).
cuda_archs_loose_intersection(SCALED_MM_3X_ARCHS "9.0;9.0a" "${CUDA_ARCHS}")
@@ -334,7 +337,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
#
# For the cutlass_scaled_mm kernels we want to build the c2x (CUTLASS 2.x)
# kernels for the remaining archs that are not already built for 3x.
- cuda_archs_loose_intersection(SCALED_MM_2X_ARCHS
+ cuda_archs_loose_intersection(SCALED_MM_2X_ARCHS
"7.5;8.0;8.6;8.7;8.9;9.0" "${CUDA_ARCHS}")
# subtract out the archs that are already built for 3x
list(REMOVE_ITEM SCALED_MM_2X_ARCHS ${SCALED_MM_3X_ARCHS})
@@ -356,6 +359,31 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
endif()
endif()
+ #
+ # 2:4 Sparse Kernels
+
+ # The 2:4 sparse kernels cutlass_scaled_sparse_mm and cutlass_compressor
+ # require CUDA 12.2 or later (and only work on Hopper, 9.0/9.0a for now).
+ if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.2 AND SCALED_MM_3X_ARCHS)
+ set(SRCS "csrc/sparse/cutlass/sparse_compressor_c3x.cu"
+ "csrc/sparse/cutlass/sparse_scaled_mm_c3x.cu")
+ set_gencode_flags_for_srcs(
+ SRCS "${SRCS}"
+ CUDA_ARCHS "${SCALED_MM_3X_ARCHS}")
+ list(APPEND VLLM_EXT_SRC "${SRCS}")
+ list(APPEND VLLM_GPU_FLAGS "-DENABLE_SPARSE_SCALED_MM_C3X=1")
+ message(STATUS "Building sparse_scaled_mm_c3x for archs: ${SCALED_MM_3X_ARCHS}")
+ else()
+ if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.2 AND SCALED_MM_3X_ARCHS)
+ message(STATUS "Not building sparse_scaled_mm_c3x kernels as CUDA Compiler version is "
+ "not >= 12.2, we recommend upgrading to CUDA 12.2 or later "
+ "if you intend on running FP8 sparse quantized models on Hopper.")
+ else()
+ message(STATUS "Not building sparse_scaled_mm_c3x as no compatible archs found "
+ "in CUDA target architectures")
+ endif()
+ endif()
+
#
# Machete kernels
@@ -442,7 +470,7 @@ define_gpu_extension_target(
SOURCES ${VLLM_EXT_SRC}
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
ARCHITECTURES ${VLLM_GPU_ARCHES}
- INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR}
+ INCLUDE_DIRECTORIES ${CUTLASS_INCLUDE_DIR};${CUTLASS_TOOLS_UTIL_INCLUDE_DIR}
USE_SABI 3
WITH_SOABI)
@@ -582,7 +610,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
- GIT_TAG 04325b6798bcc326c86fb35af62d05a9c8c8eceb
+ GIT_TAG 96266b1111111f3d11aabefaf3bacbab6a89d03c
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
diff --git a/Dockerfile b/Dockerfile
index 682f046d4b6ec..4542bc9cf0bd2 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,8 +2,8 @@
# to run the OpenAI compatible server.
# Please update any changes made here to
-# docs/source/dev/dockerfile/dockerfile.rst and
-# docs/source/assets/dev/dockerfile-stages-dependency.png
+# docs/source/contributing/dockerfile/dockerfile.md and
+# docs/source/assets/contributing/dockerfile-stages-dependency.png
ARG CUDA_VERSION=12.4.1
#################### BASE BUILD IMAGE ####################
@@ -11,6 +11,7 @@ ARG CUDA_VERSION=12.4.1
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu20.04 AS base
ARG CUDA_VERSION=12.4.1
ARG PYTHON_VERSION=3.12
+ARG TARGETPLATFORM
ENV DEBIAN_FRONTEND=noninteractive
# Install Python and other dependencies
@@ -44,12 +45,21 @@ RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
WORKDIR /workspace
# install build and runtime dependencies
+
+# arm64 (GH200) build follows the practice of "use existing pytorch" build,
+# we need to install torch and torchvision from the nightly builds first,
+# pytorch will not appear as a vLLM dependency in all of the following steps
+# after this step
+RUN --mount=type=cache,target=/root/.cache/pip \
+ if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
+ python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cu124 "torch==2.6.0.dev20241210+cu124" "torchvision==0.22.0.dev20241215"; \
+ fi
+
COPY requirements-common.txt requirements-common.txt
COPY requirements-cuda.txt requirements-cuda.txt
RUN --mount=type=cache,target=/root/.cache/pip \
python3 -m pip install -r requirements-cuda.txt
-
# cuda arch list used by torch
# can be useful for both `dev` and `test`
# explicitly set the list to avoid issues with torch 2.2
@@ -63,6 +73,7 @@ ENV VLLM_FA_CMAKE_GPU_ARCHES=${vllm_fa_cmake_gpu_arches}
#################### WHEEL BUILD IMAGE ####################
FROM base AS build
+ARG TARGETPLATFORM
# install build dependencies
COPY requirements-build.txt requirements-build.txt
@@ -134,8 +145,8 @@ COPY requirements-test.txt requirements-test.txt
COPY requirements-dev.txt requirements-dev.txt
RUN --mount=type=cache,target=/root/.cache/pip \
python3 -m pip install -r requirements-dev.txt
-
#################### DEV IMAGE ####################
+
#################### vLLM installation IMAGE ####################
# image with vLLM installed
FROM nvidia/cuda:${CUDA_VERSION}-base-ubuntu22.04 AS vllm-base
@@ -143,6 +154,7 @@ ARG CUDA_VERSION=12.4.1
ARG PYTHON_VERSION=3.12
WORKDIR /vllm-workspace
ENV DEBIAN_FRONTEND=noninteractive
+ARG TARGETPLATFORM
RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
echo "export PYTHON_VERSION_STR=${PYTHON_VERSION_STR}" >> /etc/environment
@@ -151,7 +163,7 @@ RUN PYTHON_VERSION_STR=$(echo ${PYTHON_VERSION} | sed 's/\.//g') && \
RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
&& echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \
&& apt-get update -y \
- && apt-get install -y ccache software-properties-common git curl sudo vim python3-pip \
+ && apt-get install -y ccache software-properties-common git curl wget sudo vim python3-pip \
&& apt-get install -y ffmpeg libsm6 libxext6 libgl1 \
&& add-apt-repository ppa:deadsnakes/ppa \
&& apt-get update -y \
@@ -168,18 +180,28 @@ RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
# or future versions of triton.
RUN ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/
-# install vllm wheel first, so that torch etc will be installed
+# arm64 (GH200) build follows the practice of "use existing pytorch" build,
+# we need to install torch and torchvision from the nightly builds first,
+# pytorch will not appear as a vLLM dependency in all of the following steps
+# after this step
+RUN --mount=type=cache,target=/root/.cache/pip \
+ if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
+ python3 -m pip install --index-url https://download.pytorch.org/whl/nightly/cu124 "torch==2.6.0.dev20241210+cu124" "torchvision==0.22.0.dev20241215"; \
+ fi
+
+# Install vllm wheel first, so that torch etc will be installed.
RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
--mount=type=cache,target=/root/.cache/pip \
python3 -m pip install dist/*.whl --verbose
RUN --mount=type=cache,target=/root/.cache/pip \
- . /etc/environment && \
- python3 -m pip install https://github.com/flashinfer-ai/flashinfer/releases/download/v0.1.6/flashinfer-0.1.6+cu121torch2.4-cp${PYTHON_VERSION_STR}-cp${PYTHON_VERSION_STR}-linux_x86_64.whl
+. /etc/environment && \
+if [ "$TARGETPLATFORM" != "linux/arm64" ]; then \
+ python3 -m pip install https://github.com/flashinfer-ai/flashinfer/releases/download/v0.1.6/flashinfer-0.1.6+cu121torch2.4-cp${PYTHON_VERSION_STR}-cp${PYTHON_VERSION_STR}-linux_x86_64.whl; \
+fi
COPY examples examples
#################### vLLM installation IMAGE ####################
-
#################### TEST IMAGE ####################
# image to run unit testing suite
# note that this uses vllm installed by `pip`
@@ -209,18 +231,30 @@ COPY vllm/v1 /usr/local/lib/python3.12/dist-packages/vllm/v1
RUN mkdir test_docs
RUN mv docs test_docs/
RUN mv vllm test_docs/
-
#################### TEST IMAGE ####################
#################### OPENAI API SERVER ####################
-# openai api server alternative
-FROM vllm-base AS vllm-openai
+# base openai image with additional requirements, for any subsequent openai-style images
+FROM vllm-base AS vllm-openai-base
# install additional dependencies for openai api server
RUN --mount=type=cache,target=/root/.cache/pip \
- pip install accelerate hf_transfer 'modelscope!=1.15.0' 'bitsandbytes>=0.44.0' timm==0.9.10
+ if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
+ pip install accelerate hf_transfer 'modelscope!=1.15.0' 'bitsandbytes>=0.42.0' 'timm==0.9.10' boto3 runai-model-streamer runai-model-streamer[s3]; \
+ else \
+ pip install accelerate hf_transfer 'modelscope!=1.15.0' 'bitsandbytes>=0.45.0' 'timm==0.9.10' boto3 runai-model-streamer runai-model-streamer[s3]; \
+ fi
ENV VLLM_USAGE_SOURCE production-docker-image
+# define sagemaker first, so it is not default from `docker build`
+FROM vllm-openai-base AS vllm-sagemaker
+
+COPY examples/online_serving/sagemaker-entrypoint.sh .
+RUN chmod +x sagemaker-entrypoint.sh
+ENTRYPOINT ["./sagemaker-entrypoint.sh"]
+
+FROM vllm-openai-base AS vllm-openai
+
ENTRYPOINT ["python3", "-m", "vllm.entrypoints.openai.api_server"]
#################### OPENAI API SERVER ####################
diff --git a/Dockerfile.base b/Dockerfile.base
index 1a87eb91b04b4..c064284897372 100644
--- a/Dockerfile.base
+++ b/Dockerfile.base
@@ -1,4 +1,17 @@
ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:6.3-complete
+ARG HIPBLASLT_BRANCH="4d40e36"
+ARG LEGACY_HIPBLASLT_OPTION=
+ARG RCCL_BRANCH="648a58d"
+ARG RCCL_REPO="https://github.com/ROCm/rccl"
+ARG TRITON_BRANCH="e5be006"
+ARG TRITON_REPO="https://github.com/triton-lang/triton.git"
+ARG PYTORCH_BRANCH="8d4926e"
+ARG PYTORCH_VISION_BRANCH="v0.19.1"
+ARG PYTORCH_REPO="https://github.com/pytorch/pytorch.git"
+ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
+ARG FA_BRANCH="b7d29fb"
+ARG FA_REPO="https://github.com/ROCm/flash-attention.git"
+
FROM ${BASE_IMAGE} AS base
ENV PATH=/opt/rocm/llvm/bin:$PATH
@@ -26,12 +39,12 @@ RUN apt-get update -y \
&& curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
&& python3 --version && python3 -m pip --version
-RUN pip install -U packaging cmake ninja wheel setuptools Cython
+RUN pip install -U packaging cmake ninja wheel setuptools pybind11 Cython
FROM base AS build_hipblaslt
-ARG HIPBLASLT_BRANCH="507a649"
+ARG HIPBLASLT_BRANCH
# Set to "--legacy_hipblas_direct" for ROCm<=6.2
-ARG LEGACY_HIPBLASLT_OPTION=
+ARG LEGACY_HIPBLASLT_OPTION
RUN git clone https://github.com/ROCm/hipBLASLt
RUN cd hipBLASLt \
&& git checkout ${HIPBLASLT_BRANCH} \
@@ -41,8 +54,8 @@ RUN cd hipBLASLt \
RUN mkdir -p /app/install && cp /app/hipBLASLt/build/release/*.deb /app/install
FROM base AS build_rccl
-ARG RCCL_BRANCH="dfe4a3e"
-ARG RCCL_REPO="https://github.com/ROCm/rccl"
+ARG RCCL_BRANCH
+ARG RCCL_REPO
RUN git clone ${RCCL_REPO}
RUN cd rccl \
&& git checkout ${RCCL_BRANCH} \
@@ -50,8 +63,8 @@ RUN cd rccl \
RUN mkdir -p /app/install && cp /app/rccl/build/release/*.deb /app/install
FROM base AS build_triton
-ARG TRITON_BRANCH="release/3.1.x"
-ARG TRITON_REPO="https://github.com/triton-lang/triton.git"
+ARG TRITON_BRANCH
+ARG TRITON_REPO
RUN git clone ${TRITON_REPO}
RUN cd triton \
&& git checkout ${TRITON_BRANCH} \
@@ -65,12 +78,12 @@ RUN cd /opt/rocm/share/amd_smi \
RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install
FROM base AS build_pytorch
-ARG PYTORCH_BRANCH="8bc4033"
-ARG PYTORCH_VISION_BRANCH="v0.19.1"
-ARG PYTORCH_REPO="https://github.com/pytorch/pytorch.git"
-ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
-ARG FA_BRANCH="c555642"
-ARG FA_REPO="https://github.com/ROCm/flash-attention.git"
+ARG PYTORCH_BRANCH
+ARG PYTORCH_VISION_BRANCH
+ARG PYTORCH_REPO
+ARG PYTORCH_VISION_REPO
+ARG FA_BRANCH
+ARG FA_REPO
RUN git clone ${PYTORCH_REPO} pytorch
RUN cd pytorch && git checkout ${PYTORCH_BRANCH} && \
pip install -r requirements.txt && git submodule update --init --recursive \
@@ -104,4 +117,31 @@ RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
pip install /install/*.whl
RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
- pip install /install/*.whl
\ No newline at end of file
+ pip install /install/*.whl
+
+ARG BASE_IMAGE
+ARG HIPBLASLT_BRANCH
+ARG LEGACY_HIPBLASLT_OPTION
+ARG RCCL_BRANCH
+ARG RCCL_REPO
+ARG TRITON_BRANCH
+ARG TRITON_REPO
+ARG PYTORCH_BRANCH
+ARG PYTORCH_VISION_BRANCH
+ARG PYTORCH_REPO
+ARG PYTORCH_VISION_REPO
+ARG FA_BRANCH
+ARG FA_REPO
+RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
+ && echo "HIPBLASLT_BRANCH: ${HIPBLASLT_BRANCH}" >> /app/versions.txt \
+ && echo "LEGACY_HIPBLASLT_OPTION: ${LEGACY_HIPBLASLT_OPTION}" >> /app/versions.txt \
+ && echo "RCCL_BRANCH: ${RCCL_BRANCH}" >> /app/versions.txt \
+ && echo "RCCL_REPO: ${RCCL_REPO}" >> /app/versions.txt \
+ && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
+ && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
+ && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \
+ && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \
+ && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \
+ && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt \
+ && echo "FA_BRANCH: ${FA_BRANCH}" >> /app/versions.txt \
+ && echo "FA_REPO: ${FA_REPO}" >> /app/versions.txt
\ No newline at end of file
diff --git a/Dockerfile.base_navi b/Dockerfile.base_navi
new file mode 100644
index 0000000000000..8c571ae219794
--- /dev/null
+++ b/Dockerfile.base_navi
@@ -0,0 +1,133 @@
+ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:6.3-complete
+ARG HIPBLASLT_BRANCH="4d40e36"
+ARG LEGACY_HIPBLASLT_OPTION=
+ARG RCCL_BRANCH="648a58d"
+ARG RCCL_REPO="https://github.com/ROCm/rccl"
+ARG TRITON_BRANCH="e5be006"
+ARG TRITON_REPO="https://github.com/triton-lang/triton.git"
+ARG PYTORCH_BRANCH="8d4926e"
+ARG PYTORCH_VISION_BRANCH="v0.19.1"
+ARG PYTORCH_REPO="https://github.com/pytorch/pytorch.git"
+ARG PYTORCH_VISION_REPO="https://github.com/pytorch/vision.git"
+
+FROM ${BASE_IMAGE} AS base
+
+ENV PATH=/opt/rocm/llvm/bin:$PATH
+ENV ROCM_PATH=/opt/rocm
+ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:
+ARG PYTORCH_ROCM_ARCH=gfx1100;gfx1101;gfx1200;gfx1201
+ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}
+
+ARG PYTHON_VERSION=3.12
+
+RUN mkdir -p /app
+WORKDIR /app
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install Python and other dependencies
+RUN apt-get update -y \
+ && apt-get install -y software-properties-common git curl sudo vim less \
+ && add-apt-repository ppa:deadsnakes/ppa \
+ && apt-get update -y \
+ && apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv \
+ python${PYTHON_VERSION}-lib2to3 python-is-python3 \
+ && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 \
+ && update-alternatives --set python3 /usr/bin/python${PYTHON_VERSION} \
+ && ln -sf /usr/bin/python${PYTHON_VERSION}-config /usr/bin/python3-config \
+ && curl -sS https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} \
+ && python3 --version && python3 -m pip --version
+
+RUN pip install -U packaging cmake ninja wheel setuptools Cython pybind11
+
+FROM base AS build_hipblaslt
+ARG HIPBLASLT_BRANCH
+# Set to "--legacy_hipblas_direct" for ROCm<=6.2
+ARG LEGACY_HIPBLASLT_OPTION
+RUN git clone https://github.com/ROCm/hipBLASLt
+RUN cd hipBLASLt \
+ && git checkout ${HIPBLASLT_BRANCH} \
+ && ./install.sh -d --architecture ${PYTORCH_ROCM_ARCH} ${LEGACY_HIPBLASLT_OPTION} \
+ && cd build/release \
+ && make package
+RUN mkdir -p /app/install && cp /app/hipBLASLt/build/release/*.deb /app/install
+
+FROM base AS build_rccl
+ARG RCCL_BRANCH
+ARG RCCL_REPO
+RUN git clone ${RCCL_REPO}
+RUN cd rccl \
+ && git checkout ${RCCL_BRANCH} \
+ && ./install.sh -p --amdgpu_targets ${PYTORCH_ROCM_ARCH}
+RUN mkdir -p /app/install && cp /app/rccl/build/release/*.deb /app/install
+
+FROM base AS build_triton
+ARG TRITON_BRANCH
+ARG TRITON_REPO
+RUN git clone ${TRITON_REPO}
+RUN cd triton \
+ && git checkout ${TRITON_BRANCH} \
+ && cd python \
+ && python3 setup.py bdist_wheel --dist-dir=dist
+RUN mkdir -p /app/install && cp /app/triton/python/dist/*.whl /app/install
+
+FROM base AS build_amdsmi
+RUN cd /opt/rocm/share/amd_smi \
+ && pip wheel . --wheel-dir=dist
+RUN mkdir -p /app/install && cp /opt/rocm/share/amd_smi/dist/*.whl /app/install
+
+FROM base AS build_pytorch
+ARG PYTORCH_BRANCH
+ARG PYTORCH_VISION_BRANCH
+ARG PYTORCH_REPO
+ARG PYTORCH_VISION_REPO
+RUN git clone ${PYTORCH_REPO} pytorch
+RUN cd pytorch && git checkout ${PYTORCH_BRANCH} && \
+ pip install -r requirements.txt && git submodule update --init --recursive \
+ && python3 tools/amd_build/build_amd.py \
+ && CMAKE_PREFIX_PATH=$(python3 -c 'import sys; print(sys.prefix)') python3 setup.py bdist_wheel --dist-dir=dist \
+ && pip install dist/*.whl
+RUN git clone ${PYTORCH_VISION_REPO} vision
+RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \
+ && python3 setup.py bdist_wheel --dist-dir=dist \
+ && pip install dist/*.whl
+RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \
+ && cp /app/vision/dist/*.whl /app/install
+
+FROM base AS final
+RUN --mount=type=bind,from=build_hipblaslt,src=/app/install/,target=/install \
+ dpkg -i /install/*deb \
+ && sed -i 's/, hipblaslt-dev \(.*\), hipcub-dev/, hipcub-dev/g' /var/lib/dpkg/status \
+ && sed -i 's/, hipblaslt \(.*\), hipfft/, hipfft/g' /var/lib/dpkg/status
+RUN --mount=type=bind,from=build_rccl,src=/app/install/,target=/install \
+ dpkg -i /install/*deb \
+ && sed -i 's/, rccl-dev \(.*\), rocalution/, rocalution/g' /var/lib/dpkg/status \
+ && sed -i 's/, rccl \(.*\), rocalution/, rocalution/g' /var/lib/dpkg/status
+RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
+ pip install /install/*.whl
+RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
+ pip install /install/*.whl
+RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
+ pip install /install/*.whl
+
+ARG BASE_IMAGE
+ARG HIPBLASLT_BRANCH
+ARG LEGACY_HIPBLASLT_OPTION
+ARG RCCL_BRANCH
+ARG RCCL_REPO
+ARG TRITON_BRANCH
+ARG TRITON_REPO
+ARG PYTORCH_BRANCH
+ARG PYTORCH_VISION_BRANCH
+ARG PYTORCH_REPO
+ARG PYTORCH_VISION_REPO
+RUN echo "BASE_IMAGE: ${BASE_IMAGE}" > /app/versions.txt \
+ && echo "HIPBLASLT_BRANCH: ${HIPBLASLT_BRANCH}" >> /app/versions.txt \
+ && echo "LEGACY_HIPBLASLT_OPTION: ${LEGACY_HIPBLASLT_OPTION}" >> /app/versions.txt \
+ && echo "RCCL_BRANCH: ${RCCL_BRANCH}" >> /app/versions.txt \
+ && echo "RCCL_REPO: ${RCCL_REPO}" >> /app/versions.txt \
+ && echo "TRITON_BRANCH: ${TRITON_BRANCH}" >> /app/versions.txt \
+ && echo "TRITON_REPO: ${TRITON_REPO}" >> /app/versions.txt \
+ && echo "PYTORCH_BRANCH: ${PYTORCH_BRANCH}" >> /app/versions.txt \
+ && echo "PYTORCH_VISION_BRANCH: ${PYTORCH_VISION_BRANCH}" >> /app/versions.txt \
+ && echo "PYTORCH_REPO: ${PYTORCH_REPO}" >> /app/versions.txt \
+ && echo "PYTORCH_VISION_REPO: ${PYTORCH_VISION_REPO}" >> /app/versions.txt
\ No newline at end of file
diff --git a/Dockerfile.cpu b/Dockerfile.cpu
index ebe226cf6d148..f163edc27cba8 100644
--- a/Dockerfile.cpu
+++ b/Dockerfile.cpu
@@ -26,10 +26,10 @@ RUN pip install intel_extension_for_pytorch==2.5.0
WORKDIR /workspace
+COPY requirements-build.txt requirements-build.txt
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
RUN --mount=type=cache,target=/root/.cache/pip \
- --mount=type=bind,src=requirements-build.txt,target=requirements-build.txt \
pip install --upgrade pip && \
pip install -r requirements-build.txt
@@ -37,9 +37,9 @@ FROM cpu-test-1 AS build
WORKDIR /workspace/vllm
+COPY requirements-common.txt requirements-common.txt
+COPY requirements-cpu.txt requirements-cpu.txt
RUN --mount=type=cache,target=/root/.cache/pip \
- --mount=type=bind,src=requirements-common.txt,target=requirements-common.txt \
- --mount=type=bind,src=requirements-cpu.txt,target=requirements-cpu.txt \
pip install -v -r requirements-cpu.txt
COPY . .
diff --git a/Dockerfile.neuron b/Dockerfile.neuron
index 76dbd4c04d3f3..e9cb82889decd 100644
--- a/Dockerfile.neuron
+++ b/Dockerfile.neuron
@@ -1,5 +1,6 @@
# default base image
-ARG BASE_IMAGE="public.ecr.aws/neuron/pytorch-inference-neuronx:2.1.2-neuronx-py310-sdk2.20.0-ubuntu20.04"
+# https://gallery.ecr.aws/neuron/pytorch-inference-neuronx
+ARG BASE_IMAGE="public.ecr.aws/neuron/pytorch-inference-neuronx:2.5.1-neuronx-py310-sdk2.21.0-ubuntu22.04"
FROM $BASE_IMAGE
@@ -14,16 +15,17 @@ RUN apt-get update && \
ffmpeg libsm6 libxext6 libgl1
### Mount Point ###
-# When launching the container, mount the code directory to /app
-ARG APP_MOUNT=/app
+# When launching the container, mount the code directory to /workspace
+ARG APP_MOUNT=/workspace
VOLUME [ ${APP_MOUNT} ]
WORKDIR ${APP_MOUNT}/vllm
RUN python3 -m pip install --upgrade pip
RUN python3 -m pip install --no-cache-dir fastapi ninja tokenizers pandas
-RUN python3 -m pip install sentencepiece transformers==4.36.2 -U
+RUN python3 -m pip install sentencepiece transformers==4.45.2 -U
RUN python3 -m pip install transformers-neuronx --extra-index-url=https://pip.repos.neuron.amazonaws.com -U
-RUN python3 -m pip install --pre neuronx-cc==2.15.* --extra-index-url=https://pip.repos.neuron.amazonaws.com -U
+RUN python3 -m pip install neuronx-cc==2.16.345.0 --extra-index-url=https://pip.repos.neuron.amazonaws.com -U
+RUN python3 -m pip install pytest
COPY . .
ARG GIT_REPO_CHECK=0
@@ -41,4 +43,7 @@ RUN --mount=type=bind,source=.git,target=.git \
# install development dependencies (for testing)
RUN python3 -m pip install -e tests/vllm_test_utils
+# overwrite entrypoint to run bash script
+RUN echo "import subprocess; import sys; subprocess.check_call(sys.argv[1:])" > /usr/local/bin/dockerd-entrypoint.py
+
CMD ["/bin/bash"]
diff --git a/Dockerfile.openvino b/Dockerfile.openvino
index 8bd188ffde408..32bcbfa9cc168 100644
--- a/Dockerfile.openvino
+++ b/Dockerfile.openvino
@@ -14,6 +14,7 @@ ARG GIT_REPO_CHECK=0
RUN --mount=type=bind,source=.git,target=.git \
if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
+RUN python3 -m pip install -U pip
# install build requirements
RUN PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu" python3 -m pip install -r /workspace/requirements-build.txt
# build vLLM with OpenVINO backend
diff --git a/Dockerfile.ppc64le b/Dockerfile.ppc64le
index 971248577983f..d3cd1c7b313bc 100644
--- a/Dockerfile.ppc64le
+++ b/Dockerfile.ppc64le
@@ -4,7 +4,7 @@ USER root
ENV PATH="/usr/local/cargo/bin:$PATH:/opt/conda/bin/"
-RUN apt-get update -y && apt-get install -y git wget curl vim libnuma-dev libsndfile-dev libprotobuf-dev build-essential ffmpeg libsm6 libxext6 libgl1
+RUN apt-get update -y && apt-get install -y git wget curl vim libnuma-dev libsndfile-dev libprotobuf-dev build-essential ffmpeg libsm6 libxext6 libgl1 libssl-dev
# Some packages in requirements-cpu are installed here
# IBM provides optimized packages for ppc64le processors in the open-ce project for mamba
@@ -18,9 +18,8 @@ ARG GIT_REPO_CHECK=0
RUN --mount=type=bind,source=.git,target=.git \
if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh; fi
-# These packages will be in rocketce eventually
RUN --mount=type=cache,target=/root/.cache/pip \
- pip install -v --prefer-binary --extra-index-url https://repo.fury.io/mgiessing \
+ RUSTFLAGS='-L /opt/conda/lib' pip install -v --prefer-binary --extra-index-url https://repo.fury.io/mgiessing \
'cmake>=3.26' ninja packaging 'setuptools-scm>=8' wheel jinja2 \
torch==2.3.1 \
-r requirements-cpu.txt \
diff --git a/Dockerfile.rocm b/Dockerfile.rocm
index 2334024ec65b5..13b0a764cbd7e 100644
--- a/Dockerfile.rocm
+++ b/Dockerfile.rocm
@@ -3,10 +3,13 @@ ARG REMOTE_VLLM="0"
ARG USE_CYTHON="0"
ARG BUILD_RPD="1"
ARG COMMON_WORKDIR=/app
-ARG BASE_IMAGE=rocm/vllm-dev:base_ubuntu22.04_py3.12_ROCm6.3_hipblaslt0.11_torch2.6
+ARG BASE_IMAGE=rocm/vllm-dev:base
FROM ${BASE_IMAGE} AS base
+ARG ARG_PYTORCH_ROCM_ARCH
+ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
+
# Install some basic utilities
RUN apt-get update -q -y && apt-get install -q -y \
sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev
diff --git a/README.md b/README.md
index cfeb24cbb5823..67c557bfe13a9 100644
--- a/README.md
+++ b/README.md
@@ -15,7 +15,12 @@ Easy, fast, and cheap LLM serving for everyone
---
+The first vLLM meetup in 2025 is happening on January 22nd, Wednesday, with Google Cloud in San Francisco! We will talk about vLLM's performant V1 architecture, Q1 roadmap, Google Cloud's innovation around vLLM: networking, Cloud Run, Vertex, and TPU! [Register Now](https://lu.ma/zep56hui)
+
+---
+
*Latest News* 🔥
+- [2024/12] vLLM joins [pytorch ecosystem](https://pytorch.org/blog/vllm-joins-pytorch)! Easy, Fast, and Cheap LLM Serving for Everyone!
- [2024/11] We hosted [the seventh vLLM meetup](https://lu.ma/h0qvrajz) with Snowflake! Please find the meetup slides from vLLM team [here](https://docs.google.com/presentation/d/1e3CxQBV3JsfGp30SwyvS3eM_tW-ghOhJ9PAJGK6KR54/edit?usp=sharing), and Snowflake team [here](https://docs.google.com/presentation/d/1qF3RkDAbOULwz9WK5TOltt2fE9t6uIc_hVNLFAaQX6A/edit?usp=sharing).
- [2024/10] We have just created a developer slack ([slack.vllm.ai](https://slack.vllm.ai)) focusing on coordinating contributions and discussing features. Please feel free to join us there!
- [2024/10] Ray Summit 2024 held a special track for vLLM! Please find the opening talk slides from the vLLM team [here](https://docs.google.com/presentation/d/1B_KQxpHBTRa_mDF-tR6i8rWdOU5QoTZNcEg2MKZxEHM/edit?usp=sharing). Learn more from the [talks](https://www.youtube.com/playlist?list=PLzTswPQNepXl6AQwifuwUImLPFRVpksjR) from other vLLM contributors and users!
@@ -36,7 +41,7 @@ vLLM is a fast and easy-to-use library for LLM inference and serving.
vLLM is fast with:
- State-of-the-art serving throughput
-- Efficient management of attention key and value memory with **PagedAttention**
+- Efficient management of attention key and value memory with [**PagedAttention**](https://blog.vllm.ai/2023/06/20/vllm.html)
- Continuous batching of incoming requests
- Fast model execution with CUDA/HIP graph
- Quantizations: [GPTQ](https://arxiv.org/abs/2210.17323), [AWQ](https://arxiv.org/abs/2306.00978), INT4, INT8, and FP8.
@@ -59,7 +64,7 @@ vLLM is flexible and easy to use with:
vLLM seamlessly supports most popular open-source models on HuggingFace, including:
- Transformer-like LLMs (e.g., Llama)
-- Mixture-of-Expert LLMs (e.g., Mixtral)
+- Mixture-of-Expert LLMs (e.g., Mixtral, Deepseek-V2 and V3)
- Embedding Models (e.g. E5-Mistral)
- Multi-modal LLMs (e.g., LLaVA)
@@ -76,7 +81,7 @@ pip install vllm
Visit our [documentation](https://vllm.readthedocs.io/en/latest/) to learn more.
- [Installation](https://vllm.readthedocs.io/en/latest/getting_started/installation.html)
- [Quickstart](https://vllm.readthedocs.io/en/latest/getting_started/quickstart.html)
-- [Supported Models](https://vllm.readthedocs.io/en/latest/models/supported_models.html)
+- [List of Supported Models](https://vllm.readthedocs.io/en/latest/models/supported_models.html)
## Contributing
@@ -89,28 +94,33 @@ vLLM is a community project. Our compute resources for development and testing a
-
+Cash Donations:
- a16z
+- Dropbox
+- Sequoia Capital
+- Skywork AI
+- ZhenFund
+
+Compute Resources:
- AMD
- Anyscale
- AWS
- Crusoe Cloud
- Databricks
- DeepInfra
-- Dropbox
- Google Cloud
- Lambda Lab
- Nebius
+- Novita AI
- NVIDIA
- Replicate
- Roblox
- RunPod
-- Sequoia Capital
-- Skywork AI
- Trainy
- UC Berkeley
- UC San Diego
-- ZhenFund
+
+Slack Sponsor: Anyscale
We also have an official fundraising venue through [OpenCollective](https://opencollective.com/vllm). We plan to use the fund to support the development, maintenance, and adoption of vLLM.
@@ -133,3 +143,7 @@ If you use vLLM for your research, please cite our [paper](https://arxiv.org/abs
* For coordinating contributions and development, please use Slack.
* For security disclosures, please use Github's security advisory feature.
* For collaborations and partnerships, please contact us at vllm-questions AT lists.berkeley.edu.
+
+## Media Kit
+
+* If you wish to use vLLM's logo, please refer to [our media kit repo](https://github.com/vllm-project/media-kit).
diff --git a/SECURITY.md b/SECURITY.md
index ad3f1f16ab560..de0032d26c87b 100644
--- a/SECURITY.md
+++ b/SECURITY.md
@@ -4,7 +4,7 @@
If you believe you have found a security vulnerability in vLLM, we encourage you to let us know right away. We will investigate all legitimate reports and do our best to quickly fix the problem.
-Please report security issues privately using [the vulnerability submission form](https://github.com/vllm-project/vllm/security/advisories/new).
+Please report security issues privately using [the vulnerability submission form](https://github.com/vllm-project/vllm/security/advisories/new). Reports will then be triaged by the [vulnerability management team](https://docs.vllm.ai/contributing/vulnerability_management/).
---
diff --git a/benchmarks/P3L_mling.py b/benchmarks/P3L_mling.py
new file mode 100755
index 0000000000000..0606cc0738068
--- /dev/null
+++ b/benchmarks/P3L_mling.py
@@ -0,0 +1,250 @@
+#!/usr/bin/env python3
+"""
+*MULTILINGUAL* Patch-Perplexity (P3L)
+
+This is a script that produces a realistic PPL measurement
+for the quantized KV cache system by processing a sequence of
+non-overlapping patches of the reference text. Generation of the
+consecutive symbols in each patch is governed (forced)
+by the reference text.
+
+The initial context size for the system is set by the parameter
+"--context-size".
+
+The number of output symbols to generate starting from a given
+context is set by the parameter "--sample-size". This variable also
+defines the size of the individual patch.
+
+For the N-token reference text that is split into M patches with the
+system's context size C it takes M*preload + (N-C)*generation time.
+
+Quick correctness validation tips:
+
+Running DeepSeek-V2 model
+(
+ ./vllm/examples/P3L_mling.py
+ --model=meta-llama/Llama-2-7b-chat-hf
+ --context-size=1024
+ --sample-size=512
+)
+
+should result in PPL ~ 8.42927
+
+Running DeepSeek-V2 model
+(
+ ./vllm/examples/P3L_mling.py
+ --model=meta-llama/Llama-2-7b-chat-hf
+ --context-size=1024
+ --sample-size=512
+ --patch-size=1
+ --lang-script="cmn_Hant"
+)
+should result in PPL ~ 2.67962
+
+The multi-linguality is implemented through the additional
+key "--lang-script", which defaults to English in Latin
+scripture ("eng_Latn").
+
+Please refer to
+
+https://confluence.amd.com/display/MLSE/Multi-Lingual+P3L+Test
+
+for the complete set of possible language-scripture choices.
+
+
+"""
+
+import argparse
+import dataclasses
+import datetime
+import json
+import math
+import os
+
+import pandas
+from huggingface_hub import hf_hub_download
+
+from vllm import LLM, SamplingParams
+from vllm.engine.arg_utils import EngineArgs
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def get_wikitext2_text(tokenizer):
+ hf_hub_download(repo_id='alexei-v-ivanov-amd/wiki',
+ repo_type="dataset",
+ filename='wiki.test.raw',
+ local_dir='./')
+ with open('./wiki.test.raw') as f:
+ test_text = "\n".join(line.strip() for line in f)
+ test_enc = tokenizer(test_text)
+
+ os.remove('./wiki.test.raw')
+
+ return test_enc, test_text
+
+
+def get_flores_plus_text(tokenizer, lng_scrpt):
+ hf_hub_download(repo_id='alexei-v-ivanov-amd/flores_plus',
+ repo_type="dataset",
+ filename=lng_scrpt + '.parquet',
+ local_dir='./')
+
+ df = pandas.read_parquet('./' + lng_scrpt + '.parquet')
+ test_text = "\n\n".join(line.strip() for line in df['text'])
+ test_enc = tokenizer(test_text)
+
+ os.remove('./' + lng_scrpt + '.parquet')
+
+ return test_enc, test_text
+
+
+def vllm_init(args):
+ engine_args = EngineArgs.from_cli_args(args)
+ llm = LLM(**dataclasses.asdict(engine_args))
+
+ sampling_params = SamplingParams(n=1,
+ temperature=0.0,
+ top_p=1,
+ ignore_eos=True,
+ ppl_measurement=True,
+ future_context=[],
+ prompt_logprobs=1,
+ logprobs=1,
+ presence_penalty=0.0)
+
+ return llm, sampling_params
+
+
+def vllm_predict(CONT, llm, sampl_par):
+ result = llm.generate(prompt_token_ids=CONT, sampling_params=sampl_par)
+ return result
+
+
+def main(args: argparse.Namespace):
+
+ MESSAGE = f"Initialising @ {datetime.datetime.now()}"
+ logger.info(MESSAGE)
+ print(MESSAGE)
+ my_ppl = 0.0
+
+ logger.info("Initializing the engine.")
+ my_llm, my_sampl_par = vllm_init(args)
+ my_tokenizer = my_llm.llm_engine.tokenizer.tokenizer
+ logger.info(my_sampl_par)
+ logger.info("Initialized the engine.")
+
+ my_n_samples = args.sample_size
+ my_lang_script = args.lang_script
+
+ if (args.context_size+my_n_samples) > \
+ my_llm.llm_engine.model_config.max_model_len:
+ MESSAGE = ("" \
+ "Error! The total number of tokens:\n" \
+ f" prefix ({args.context_size}) + " \
+ f"to be generated ({my_n_samples})" \
+ f" can't be bigger than the model limit " \
+ f"({my_llm.llm_engine.model_config.max_model_len}).")
+ logger.info(MESSAGE)
+ print(MESSAGE)
+ return
+
+ my_test_enc, my_test_text = get_flores_plus_text(my_tokenizer,
+ my_lang_script)
+
+ logger.info("Loaded the test data.")
+
+ my_n_patches = math.ceil(
+ (len(my_test_enc['input_ids']) - args.context_size - 1) / my_n_samples)
+ if args.patch_size is not None:
+ my_n_patches = args.patch_size
+
+ num_tokens_generated = 0
+ starting_time = datetime.datetime.now()
+ MESSAGE = (f"Starting generation @ {starting_time}\n" \
+ " Have the test sample of "
+ f"{len(my_test_enc['input_ids'])} tokens" \
+ f" will try to process {my_n_patches} patche(s)," \
+ f" generating {my_n_samples} tokens in each patch" \
+ f" from the initial context of {args.context_size} tokens.")
+
+ logger.info(MESSAGE)
+ print(MESSAGE)
+ for c in range(my_n_patches):
+ CONTEXT = []
+ my_sampl_par.future_context = []
+ CONTEXT.append(
+ my_test_enc['input_ids'][c * my_n_samples:c * my_n_samples +
+ args.context_size])
+ upper_boundary = min((c + 1) * my_n_samples + args.context_size,
+ len(my_test_enc['input_ids']))
+ my_sampl_par.future_context.append(
+ my_test_enc['input_ids'][c * my_n_samples +
+ args.context_size:upper_boundary])
+ my_sampl_par.max_tokens = len(my_sampl_par.future_context[0])
+ my_sampl_par.cntr = c
+ LOGPROBS = vllm_predict(CONTEXT, my_llm, my_sampl_par)
+ num_tokens_generated += len(LOGPROBS[0].outputs[0].token_ids)
+ if (num_tokens_generated < my_n_samples):
+ MESSAGE = (f"Warning: The number of generated tokens is" \
+ f"less than requested ({num_tokens_generated}" \
+ f" < {my_n_samples}).")
+ logger.info(MESSAGE)
+ print(MESSAGE)
+ my_ppl -= LOGPROBS[0].outputs[0].cumulative_logprob
+ MESSAGE = (f"Iteration {c+1} of {my_n_patches} Intermediate" \
+ "Estimates:\n" \
+ f"\tCross-entropy_intermediate={my_ppl/num_tokens_generated}\n" \
+ f"\tPerplexity_intermediate=" \
+ f"{math.exp(my_ppl/num_tokens_generated)}")
+
+ logger.info(MESSAGE)
+ print(MESSAGE)
+ ending_time = datetime.datetime.now()
+ MESSAGE = (f"Done @ {ending_time} after processing for" \
+ f" {ending_time-starting_time}" \
+ f" generated {num_tokens_generated} tokens.")
+
+ logger.info(MESSAGE)
+ print(MESSAGE)
+
+ MESSAGE = (f"\tIntegral Cross-Entropy={my_ppl}\n\tAverage Cross-Entropy=" \
+ f"{my_ppl/num_tokens_generated}" \
+ f"\n\tPPL={math.exp(my_ppl/num_tokens_generated)}")
+
+ if args.output_json:
+ results = {
+ "integral_cross_entropy": my_ppl,
+ "average_cross_entropy": my_ppl / num_tokens_generated,
+ "ppl": math.exp(my_ppl / num_tokens_generated),
+ }
+ with open(args.output_json, "w") as f:
+ json.dump(results, f, indent=4)
+
+ logger.info(MESSAGE)
+ print(MESSAGE)
+ return
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(
+ description='Measure the PPPL (P3L) score of a given model.')
+ parser.add_argument(
+ '--data',
+ type=str,
+ default='./wikitext/wikitext-2-v1/test-00000-of-00001.parquet')
+ parser.add_argument('--context-size', type=int, default=4096)
+ parser.add_argument('--sample-size', type=int, default=512)
+ parser.add_argument('--patch-size', type=int, default=None)
+ parser.add_argument('--lang-script', type=str, default="eng_Latn")
+ parser.add_argument(
+ '--output-json',
+ type=str,
+ default=None,
+ help='Path to save the latency results in JSON format.')
+
+ parser = EngineArgs.add_cli_args(parser)
+ args = parser.parse_args()
+
+ main(args)
diff --git a/benchmarks/backend_request_func.py b/benchmarks/backend_request_func.py
index c3fed56e8a956..9d71e4ecc4a37 100644
--- a/benchmarks/backend_request_func.py
+++ b/benchmarks/backend_request_func.py
@@ -24,6 +24,7 @@ class RequestFuncInput:
model: str
best_of: int = 1
logprobs: Optional[int] = None
+ extra_body: Optional[dict] = None
multi_modal_content: Optional[dict] = None
ignore_eos: bool = False
@@ -36,6 +37,7 @@ class RequestFuncOutput:
ttft: float = 0.0 # Time to first token
itl: List[float] = field(
default_factory=list) # List of inter-token latencies
+ tpot: float = 0.0 # avg next-token latencies
prompt_len: int = 0
error: str = ""
@@ -242,6 +244,8 @@ async def async_request_openai_completions(
"stream": True,
"ignore_eos": request_func_input.ignore_eos,
}
+ if request_func_input.extra_body:
+ payload.update(request_func_input.extra_body)
headers = {
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"
}
@@ -336,6 +340,8 @@ async def async_request_openai_chat_completions(
"stream": True,
"ignore_eos": request_func_input.ignore_eos,
}
+ if request_func_input.extra_body:
+ payload.update(request_func_input.extra_body)
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
@@ -411,14 +417,35 @@ def get_model(pretrained_model_name_or_path: str) -> str:
def get_tokenizer(
- pretrained_model_name_or_path: str, trust_remote_code: bool
+ pretrained_model_name_or_path: str,
+ tokenizer_mode: str = "auto",
+ trust_remote_code: bool = False,
+ **kwargs,
) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
if pretrained_model_name_or_path is not None and not os.path.exists(
pretrained_model_name_or_path):
pretrained_model_name_or_path = get_model(
pretrained_model_name_or_path)
- return AutoTokenizer.from_pretrained(pretrained_model_name_or_path,
- trust_remote_code=trust_remote_code)
+ if tokenizer_mode == "slow":
+ if kwargs.get("use_fast", False):
+ raise ValueError(
+ "Cannot use the fast tokenizer in slow tokenizer mode.")
+ kwargs["use_fast"] = False
+ if tokenizer_mode == "mistral":
+ try:
+ from vllm.transformers_utils.tokenizer import MistralTokenizer
+ except ImportError as e:
+ raise ImportError("MistralTokenizer requires vllm package.\n"
+ "Please install it with `pip install vllm` "
+ "to use mistral tokenizer mode.") from e
+ return MistralTokenizer.from_pretrained(
+ str(pretrained_model_name_or_path))
+ else:
+ return AutoTokenizer.from_pretrained(
+ pretrained_model_name_or_path,
+ trust_remote_code=trust_remote_code,
+ **kwargs,
+ )
ASYNC_REQUEST_FUNCS = {
diff --git a/benchmarks/benchmark_guided.py b/benchmarks/benchmark_guided.py
new file mode 100644
index 0000000000000..1a0e62598bfcb
--- /dev/null
+++ b/benchmarks/benchmark_guided.py
@@ -0,0 +1,494 @@
+"""Benchmark guided decoding throughput."""
+import argparse
+import dataclasses
+import json
+import os
+import random
+import time
+from typing import List
+
+import datasets
+import pandas as pd
+import uvloop
+from transformers import AutoTokenizer, PreTrainedTokenizerBase
+
+from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
+from vllm.entrypoints.openai.api_server import (
+ build_async_engine_client_from_engine_args)
+from vllm.sampling_params import GuidedDecodingParams
+from vllm.utils import FlexibleArgumentParser, merge_async_iterators
+
+
+@dataclasses.dataclass
+class SampleRequest:
+ """A class representing a single inference request for benchmarking.
+
+ Attributes:
+ prompt: The input text prompt for the model.
+ multi_modal_data: Optional dictionary containing multi-modal data (e.g.
+ images).
+ prompt_len: The length of the prompt in tokens.
+ expected_output_len: The expected length of the output in tokens.
+ """
+ prompt: str
+ prompt_len: int
+ expected_output_len: int
+ schema: dict
+ structure_type: str = 'json'
+ completion: str = None
+
+
+def run_vllm(requests: List[SampleRequest],
+ engine_args: EngineArgs,
+ n: int,
+ guided_decoding_rate: float = 1.0,
+ warmup: bool = False) -> float:
+ from vllm import LLM, SamplingParams
+ llm = LLM(**vars(engine_args))
+
+ # Add the requests to the engine.
+ prompts: List[str] = []
+ sampling_params: List[SamplingParams] = []
+ # create a list containing random selected true or false
+ guided_decoding_req_idx = random.sample(
+ range(len(requests)), int(len(requests) * guided_decoding_rate))
+
+ if warmup:
+ print(">>>>> Running warmup prompt, for the first 5")
+ # We setup the first 5 requests to warmup FSM
+ # if using xgrammar dataset, we will skip warmup
+ warmup_requests = requests[:5]
+ for i, request in enumerate(warmup_requests):
+ prompts.append(request.prompt)
+ sampling_params.append(
+ SamplingParams(
+ n=n,
+ temperature=1.0,
+ top_p=1.0,
+ ignore_eos=True,
+ max_tokens=request.expected_output_len,
+ guided_decoding=GuidedDecodingParams(json=request.schema)
+ if guided_decoding_rate > 0 else None,
+ ))
+ llm.generate(prompts, sampling_params, use_tqdm=False)
+
+ print(">>>>> Benchmark started...")
+ prompts = []
+ sampling_params = []
+ for i, request in enumerate(requests):
+ prompts.append(request.prompt)
+ sampling_params.append(
+ SamplingParams(
+ n=n,
+ temperature=1.0,
+ top_p=1.0,
+ ignore_eos=True,
+ max_tokens=request.expected_output_len,
+ guided_decoding=GuidedDecodingParams(
+ **{request.structure_type: request.schema})
+ if i in guided_decoding_req_idx else None,
+ ))
+
+ start = time.perf_counter()
+ outputs = llm.generate(prompts, sampling_params, use_tqdm=False)
+ ret = []
+ for output, request in zip(outputs, requests):
+ generated_text = output.outputs[0].text
+ ret.append({
+ "generated": generated_text,
+ "expected": request.completion
+ })
+ end = time.perf_counter()
+ return end - start, ret
+
+
+async def run_vllm_async(
+ requests: List[SampleRequest],
+ engine_args: AsyncEngineArgs,
+ n: int,
+ guided_decoding_rate: float = 1.0,
+ warmup: bool = False,
+ disable_frontend_multiprocessing: bool = False) -> float:
+ from vllm import SamplingParams
+
+ async with build_async_engine_client_from_engine_args(
+ engine_args, disable_frontend_multiprocessing) as llm:
+
+ # Add the requests to the engine.
+ prompts: List[str] = []
+ sampling_params: List[SamplingParams] = []
+ guided_decoding_req_idx = random.sample(
+ range(len(requests)), int(len(requests) * guided_decoding_rate))
+
+ if warmup:
+ print(">>>>>> Running warmup prompt, for the first 5")
+ # We setup the first 5 requests to warmup FSM
+ # if using xgrammar dataset, we will skip warmup
+ warmup_requests = requests[:5]
+ for i, request in enumerate(warmup_requests):
+ prompts.append(request.prompt)
+ sampling_params.append(
+ SamplingParams(
+ n=n,
+ temperature=1.0,
+ top_p=1.0,
+ ignore_eos=True,
+ max_tokens=request.expected_output_len,
+ guided_decoding=GuidedDecodingParams(
+ json=request.schema)
+ if guided_decoding_rate > 0 else None,
+ ))
+ generators = []
+ for i, (prompt, sp) in enumerate(zip(prompts, sampling_params)):
+ generator = llm.generate(prompt, sp, request_id=f"test{i}")
+ generators.append(generator)
+ all_gens = merge_async_iterators(*generators)
+ async for i, res in all_gens:
+ pass
+
+ print(">>>>> Benchmark started...")
+ prompts = []
+ sampling_params = []
+ for i, request in enumerate(requests):
+ prompts.append(request.prompt)
+ sampling_params.append(
+ SamplingParams(
+ n=n,
+ temperature=1.0,
+ top_p=1.0,
+ ignore_eos=True,
+ max_tokens=request.expected_output_len,
+ guided_decoding=GuidedDecodingParams(json=request.schema)
+ if i in guided_decoding_req_idx else None,
+ ))
+
+ generators = []
+ start_time = []
+ latencies = []
+ start = time.perf_counter()
+ for i, (prompt, sp) in enumerate(zip(prompts, sampling_params)):
+ generator = llm.generate(prompt, sp, request_id=f"test{i}")
+ generators.append(generator)
+ start_time.append(time.perf_counter())
+ latencies.append([])
+ all_gens = merge_async_iterators(*generators)
+ generated_texts = [''] * len(requests)
+ async for i, res in all_gens:
+ generated_texts[i] = res.outputs[0].text
+ lat = time.perf_counter() - start_time[i]
+ latencies[i].append(lat)
+ ret = [{
+ 'generated': gt,
+ 'expected': req.completion
+ } for gt, req in zip(generated_texts, requests)]
+ end = time.perf_counter()
+ first_latency = pd.Series([lat[0] * 1000 for lat in latencies])
+ next_latency = pd.Series([(lat[-1] - lat[0]) / len(lat[1:]) * 1000
+ for lat in latencies])
+ return end - start, ret, (first_latency, next_latency)
+
+
+def sample_requests(tokenizer: PreTrainedTokenizerBase,
+ args: argparse.Namespace) -> List[SampleRequest]:
+ if args.dataset == 'json':
+ if args.json_schema_path is None:
+ dir_path = os.path.dirname(os.path.realpath(__file__))
+ args.json_schema_path = os.path.join(dir_path,
+ "structured_schemas",
+ "structured_schema_1.json")
+ with open(args.json_schema_path) as f:
+ schema = json.load(f)
+ prompt = f"Generate an example of a user profile given the following schema: {json.dumps(schema)}" # noqa: E501
+ input_len = len(tokenizer(prompt).input_ids)
+ print(f"Input length of the prompt: {input_len} tokens")
+ requests = [
+ SampleRequest(prompt=prompt,
+ prompt_len=input_len,
+ expected_output_len=args.output_len,
+ schema=schema,
+ structure_type=args.structure_type)
+ for _ in range(args.num_prompts)
+ ]
+
+ elif args.dataset == "grammar":
+ schema = """
+ ?start: select_statement
+
+ ?select_statement: "SELECT " column_list " FROM " table_name
+
+ ?column_list: column_name ("," column_name)*
+
+ ?table_name: identifier
+
+ ?column_name: identifier
+
+ ?identifier: /[a-zA-Z_][a-zA-Z0-9_]*/
+ """
+ prompt = "Generate an SQL query to show the 'username' \
+ and 'email' from the 'users' table."
+
+ input_len = len(tokenizer(prompt).input_ids)
+ print(f"Input length of the prompt: {input_len} tokens")
+ requests = [
+ SampleRequest(prompt=prompt,
+ prompt_len=input_len,
+ expected_output_len=args.output_len,
+ schema=schema,
+ structure_type=args.structure_type)
+ for _ in range(args.num_prompts)
+ ]
+
+ elif args.dataset == "regex":
+ regex = r"\w+@\w+\.com\n"
+ args.regex = regex
+ prompt = "Generate an email address for Alan Turing, \
+ who works in Enigma. End in .com and new line. \
+ Example result: alan.turing@enigma.com\n"
+
+ input_len = len(tokenizer(prompt).input_ids)
+ print(f"Input length of the prompt: {input_len} tokens")
+ requests = [
+ SampleRequest(prompt=prompt,
+ prompt_len=input_len,
+ expected_output_len=args.output_len,
+ schema=regex,
+ structure_type=args.structure_type)
+ for _ in range(args.num_prompts)
+ ]
+
+ elif args.dataset == "choice":
+ choice = ["Positive", "Negative"]
+ args.choice = choice
+ prompt = "Classify this sentiment: vLLM is wonderful!"
+ input_len = len(tokenizer(prompt).input_ids)
+ print(f"Input length of the prompt: {input_len} tokens")
+ requests = [
+ SampleRequest(prompt=prompt,
+ prompt_len=input_len,
+ expected_output_len=args.output_len,
+ schema=choice,
+ structure_type=args.structure_type)
+ for _ in range(args.num_prompts)
+ ]
+
+ elif args.dataset == "xgrammar_bench":
+ args.warmup = False
+ requests: List[SampleRequest] = []
+ dataset = datasets.load_dataset("NousResearch/json-mode-eval",
+ split="train")
+ print(f"dataset has {len(dataset)} entries")
+ len_dataset = len(dataset)
+ for data_point_idx in range(args.num_prompts):
+ idx = data_point_idx
+ while idx >= len_dataset:
+ idx -= len_dataset
+ schema = dataset["schema"][idx]
+ prompt = tokenizer.apply_chat_template(dataset["prompt"][idx],
+ tokenize=False)
+ input_len = len(tokenizer(prompt).input_ids)
+ completion = dataset["completion"][idx]
+
+ requests.append(
+ SampleRequest(prompt=prompt,
+ prompt_len=input_len,
+ expected_output_len=args.output_len,
+ schema=schema,
+ completion=completion))
+
+ return requests
+
+
+def evaluate(ret, args):
+
+ def _eval_correctness_json(expected, actual):
+ # extract json string from string using regex
+ import re
+ actual = actual.replace('\n', '').replace(' ', '').strip()
+ try:
+ actual = re.search(r'\{.*\}', actual).group()
+ actual = json.loads(actual)
+ except Exception:
+ return False
+
+ return True
+
+ def _eval_correctness_choice(expected, actual):
+ return actual in args.choice
+
+ def _eval_correctness_regex(expected, actual):
+ import re
+ return re.match(args.regex, actual) is not None
+
+ def _eval_correctness(expected, actual):
+ if args.structure_type == 'json':
+ return _eval_correctness_json(expected, actual)
+ elif args.structure_type == 'regex':
+ return _eval_correctness_regex(expected, actual)
+ elif args.structure_type == 'choice':
+ return _eval_correctness_choice(expected, actual)
+ else:
+ return None
+
+ scores = []
+ for res in ret:
+ score = _eval_correctness(res['expected'], res['generated'])
+ res['correctness'] = score
+ scores.append(score)
+
+ not_none_scores = [score for score in scores if score is not None]
+
+ return (sum(not_none_scores) / len(not_none_scores) *
+ 100) if len(not_none_scores) > 0 else None
+
+
+def main(args: argparse.Namespace):
+ print(args)
+ random.seed(args.seed)
+
+ # async engine is working for 'regex', 'choice' and 'grammar'
+ if args.dataset == 'grammar':
+ args.structure_type = 'grammar'
+ args.async_engine = False
+ elif args.dataset == 'regex':
+ args.structure_type = 'regex'
+ args.async_engine = False
+ elif args.dataset == 'choice':
+ args.structure_type = 'choice'
+ args.async_engine = False
+ else:
+ args.structure_type = 'json'
+
+ if args.no_guided_decoding:
+ args.guided_decoding_ratio = 0
+ if args.save_results:
+ result_file_name = f'{args.guided_decoding_ratio}guided'
+ result_file_name += f"_{args.model.split('/')[-1]}"
+ result_file_name += f"_{args.dataset}"
+ result_file_name += f"_{args.num_prompts}"
+ result_file_name += f"_out{args.output_len}"
+ result_file_name += f"_async{args.async_engine}"
+ result_file_name += f"_warmup{args.warmup}"
+ result_file_name += f"_chunkedprefill{args.enable_chunked_prefill}"
+ result_file_name += ".txt"
+ else:
+ result_file_name = None
+
+ # Synthesize a prompt with the given input length.
+ tokenizer = AutoTokenizer.from_pretrained(
+ args.tokenizer, trust_remote_code=args.trust_remote_code)
+ requests = sample_requests(tokenizer, args)
+
+ if args.async_engine:
+ engine_args = AsyncEngineArgs.from_cli_args(args)
+ elapsed_time, ret, (first_latency, next_latency) = uvloop.run(
+ run_vllm_async(requests, engine_args, args.n,
+ args.guided_decoding_ratio, args.warmup,
+ args.disable_frontend_multiprocessing))
+ else:
+ engine_args = EngineArgs.from_cli_args(args)
+ elapsed_time, ret = run_vllm(requests, engine_args, args.n,
+ args.guided_decoding_ratio, args.warmup)
+ first_latency, next_latency = None, None
+
+ score = evaluate(ret, args)
+ total_num_tokens = sum(request.prompt_len + request.expected_output_len
+ for request in requests)
+ total_output_tokens = sum(request.expected_output_len
+ for request in requests)
+ if first_latency is not None:
+ latency_breakdown = "\nFirst token latency(msecs):\n"
+ latency_breakdown += f"{first_latency.describe()}"
+ latency_breakdown += "\nNext token latency(msecs):\n"
+ latency_breakdown += f"{next_latency.describe()}"
+ print(
+ f"Throughput: {len(requests) / elapsed_time:.2f} requests/s, "
+ f"{total_num_tokens / elapsed_time:.2f} total tokens/s, "
+ f"{total_output_tokens / elapsed_time:.2f} output tokens/s",
+ f"Correct rate is {score} %",
+ f"{latency_breakdown if first_latency is not None else ''}")
+
+ # Output JSON results if specified
+ if args.output_json or result_file_name:
+ results = {
+ "elapsed_time": elapsed_time,
+ "num_requests": len(requests),
+ "total_num_tokens": total_num_tokens,
+ "total_output_tokens": total_output_tokens,
+ "requests_per_second": len(requests) / elapsed_time,
+ "tokens_per_second": f"{total_num_tokens / elapsed_time:.2f}",
+ "output_tokens_per_second":
+ f"{total_output_tokens / elapsed_time:.2f}",
+ "correct_rate(%)": score
+ }
+ results = {"outputs": ret, **results}
+ if first_latency is not None:
+ results["first_token_latency(msecs)"] = first_latency.describe(
+ ).to_dict()
+ results["next_token_latency(msecs)"] = next_latency.describe(
+ ).to_dict()
+ if args.output_json:
+ with open(args.output_json, "w") as f:
+ json.dump(results, f, indent=4)
+ elif result_file_name:
+ with open(result_file_name, "w") as f:
+ json.dump(results, f, indent=4)
+
+
+if __name__ == "__main__":
+ parser = FlexibleArgumentParser(description="Benchmark guided decoding.")
+ parser = AsyncEngineArgs.add_cli_args(parser)
+
+ parser.add_argument("--output-len",
+ type=int,
+ default=512,
+ help="Output length for each request. Overrides the "
+ "output length from the dataset.")
+ parser.add_argument(
+ "--dataset",
+ default='json',
+ choices=['json', 'grammar', 'regex', 'choice', 'xgrammar_bench'])
+ parser.add_argument("--json_schema_path",
+ type=str,
+ default=None,
+ help="Path to json schema.")
+ parser.add_argument("--n",
+ type=int,
+ default=1,
+ help="Number of generated sequences per prompt.")
+ parser.add_argument("--num-prompts",
+ type=int,
+ default=10,
+ help="Number of prompts to process.")
+ parser.add_argument(
+ '--output-json',
+ type=str,
+ default=None,
+ help='Path to save the throughput results in JSON format.')
+ parser.add_argument("--async-engine",
+ action='store_true',
+ default=False,
+ help="Use vLLM async engine rather than LLM class.")
+ parser.add_argument("--no-guided-decoding",
+ action='store_true',
+ default=False,
+ help="Whether to disable JSON decoding or not.")
+ parser.add_argument("--guided-decoding-ratio",
+ type=float,
+ default=1.0,
+ help="Ratio of Guided Decoding requests")
+ parser.add_argument("--disable-frontend-multiprocessing",
+ action='store_true',
+ default=False,
+ help="Disable decoupled async engine frontend.")
+ parser.add_argument("--warmup",
+ action="store_true",
+ default=False,
+ help="Run warmup prompts before benchmark.")
+ parser.add_argument("--save-results",
+ action="store_true",
+ default=False,
+ help="save output results.")
+ args = parser.parse_args()
+ if args.tokenizer is None:
+ args.tokenizer = args.model
+ main(args)
diff --git a/benchmarks/benchmark_latency.py b/benchmarks/benchmark_latency.py
index 0a14aedd5feba..77c4f6aa927e4 100644
--- a/benchmarks/benchmark_latency.py
+++ b/benchmarks/benchmark_latency.py
@@ -13,6 +13,7 @@
from vllm import LLM, SamplingParams
from vllm.engine.arg_utils import EngineArgs
from vllm.inputs import PromptType
+from vllm.sampling_params import BeamSearchParams
from vllm.utils import FlexibleArgumentParser
@@ -40,6 +41,20 @@ def main(args: argparse.Namespace):
"prompt_token_ids": batch
} for batch in dummy_prompt_token_ids.tolist()]
+ def llm_generate():
+ if not args.use_beam_search:
+ llm.generate(dummy_prompts,
+ sampling_params=sampling_params,
+ use_tqdm=False)
+ else:
+ llm.beam_search(
+ dummy_prompts,
+ BeamSearchParams(
+ beam_width=args.n,
+ max_tokens=args.output_len,
+ ignore_eos=True,
+ ))
+
def run_to_completion(profile_dir: Optional[str] = None):
if profile_dir:
with torch.profiler.profile(
@@ -49,15 +64,11 @@ def run_to_completion(profile_dir: Optional[str] = None):
],
on_trace_ready=torch.profiler.tensorboard_trace_handler(
str(profile_dir))) as p:
- llm.generate(dummy_prompts,
- sampling_params=sampling_params,
- use_tqdm=False)
- print(p.key_averages())
+ llm_generate()
+ print(p.key_averages().table(sort_by="self_cuda_time_total"))
else:
start_time = time.perf_counter()
- llm.generate(dummy_prompts,
- sampling_params=sampling_params,
- use_tqdm=False)
+ llm_generate()
end_time = time.perf_counter()
latency = end_time - start_time
return latency
diff --git a/benchmarks/benchmark_long_document_qa_throughput.py b/benchmarks/benchmark_long_document_qa_throughput.py
new file mode 100644
index 0000000000000..0b8fba38156f1
--- /dev/null
+++ b/benchmarks/benchmark_long_document_qa_throughput.py
@@ -0,0 +1,183 @@
+"""
+Offline benchmark to test the long document QA throughput.
+
+Example usage:
+ # This workload samples 8 different prompts with a default input
+ # length of 20000 tokens, then replicates each prompt 2 times
+ # in random order.
+ python benchmark_long_document_qa_throughput.py \
+ --model meta-llama/Llama-2-7b-chat-hf \
+ --enable-prefix-caching \
+ --num-documents 8 \
+ --repeat-count 2
+
+Commandline arguments:
+ --num-documents: The number of documents to sample prompts from.
+
+ --document-length: The length of each document in tokens.
+ (Optional, default: 20000)
+
+ --output-len: The number of tokens to generate for each prompt.
+ (Optional, default: 10)
+
+ --repeat-count: The number of times to repeat each prompt.
+ (Optional, default: 2)
+
+ --repeat-mode: The mode to repeat prompts. The supported modes are:
+ - 'random': shuffle the prompts randomly. (Default)
+ - 'tile': the entire prompt list is repeated in sequence. (Potentially
+ lowest cache hit)
+ - 'interleave': each prompt is repeated consecutively before
+ moving to the next element. (Highest cache hit)
+
+ --shuffle-seed: Random seed when the repeat mode is "random".
+ (Optional, default: 0)
+
+In the meantime, it also supports all the vLLM engine args to initialize the
+LLM engine. You can refer to the `vllm.engine.arg_utils.EngineArgs` for more
+details.
+"""
+
+import dataclasses
+import random
+import time
+
+from vllm import LLM, SamplingParams
+from vllm.engine.arg_utils import EngineArgs
+from vllm.utils import FlexibleArgumentParser
+
+
+def test_long_document_qa(llm=None, sampling_params=None, prompts=None):
+ """
+ Test long document QA with the given prompts and sampling parameters.
+ Print the time spent in processing all the prompts.
+
+ Args:
+ llm: The language model used for generating responses.
+ sampling_params: Sampling parameter used to generate the response.
+ prompts: A list of prompt strings to be processed by the LLM.
+ """
+ start_time = time.time()
+ llm.generate(prompts, sampling_params=sampling_params)
+ end_time = time.time()
+ print(f"Time to execute all requests: {end_time - start_time:.4f} secs")
+
+
+def repeat_prompts(prompts, repeat_count, mode: str):
+ """
+ Repeat each prompt in the list for a specified number of times.
+ The order of prompts in the output list depends on the mode.
+
+ Args:
+ prompts: A list of prompts to be repeated.
+ repeat_count: The number of times each prompt is repeated.
+ mode: The mode of repetition. Supported modes are:
+ - 'random': Shuffle the prompts randomly after repetition.
+ - 'tile': Repeat the entire prompt list in sequence.
+ Example: [1, 2, 3] -> [1, 2, 3, 1, 2, 3].
+ - 'interleave': Repeat each prompt consecutively before moving to
+ the next. Example: [1, 2, 3] -> [1, 1, 2, 2, 3, 3].
+
+ Returns:
+ A list of repeated prompts in the specified order.
+
+ Raises:
+ ValueError: If an invalid mode is provided.
+ """
+ print("Repeat mode: ", mode)
+ if mode == 'random':
+ repeated_prompts = prompts * repeat_count
+ random.shuffle(repeated_prompts)
+ return repeated_prompts
+ elif mode == 'tile':
+ return prompts * repeat_count
+ elif mode == 'interleave':
+ repeated_prompts = []
+ for prompt in prompts:
+ repeated_prompts.extend([prompt] * repeat_count)
+ return repeated_prompts
+ else:
+ raise ValueError(f"Invalid mode: {mode}, only support "
+ "'random', 'tile', 'interleave'")
+
+
+def main(args):
+ random.seed(args.shuffle_seed)
+
+ # Prepare the prompts:
+ # we append the document id at the beginning to avoid any of the document
+ # being the prefix of other documents
+ prompts = [
+ str(i) + ' '.join(['hi'] * args.document_length)
+ for i in range(args.num_documents)
+ ]
+
+ prompts = repeat_prompts(prompts, args.repeat_count, mode=args.repeat_mode)
+
+ warmup_prompts = [
+ "This is warm up request " + str(i) + \
+ ' '.join(['hi'] * args.document_length)
+ for i in range(args.num_documents)]
+
+ # Create the LLM engine
+ engine_args = EngineArgs.from_cli_args(args)
+ llm = LLM(**dataclasses.asdict(engine_args))
+ sampling_params = SamplingParams(temperature=0, max_tokens=args.output_len)
+
+ print("------warm up------")
+ test_long_document_qa(
+ llm=llm,
+ prompts=warmup_prompts,
+ sampling_params=sampling_params,
+ )
+
+ print("------start generating------")
+ test_long_document_qa(
+ llm=llm,
+ prompts=prompts,
+ sampling_params=sampling_params,
+ )
+
+
+if __name__ == "__main__":
+ parser = FlexibleArgumentParser(
+ description=
+ 'Benchmark the performance with or without automatic prefix caching.')
+
+ parser.add_argument(
+ '--document-length',
+ type=int,
+ # Roughly the number of tokens for a system paper,
+ # excluding images
+ default=20000,
+ help='Range of input lengths for sampling prompts,'
+ 'specified as "min:max" (e.g., "128:256").')
+
+ parser.add_argument('--num-documents',
+ type=int,
+ default=8,
+ help='Range of input lengths for sampling prompts,'
+ 'specified as "min:max" (e.g., "128:256").')
+
+ parser.add_argument('--output-len', type=int, default=10)
+
+ parser.add_argument('--repeat-count',
+ type=int,
+ default=2,
+ help='Number of times to repeat each prompt')
+
+ parser.add_argument("--repeat-mode",
+ type=str,
+ default='random',
+ help='The mode to repeat prompts. The supported '
+ 'modes are "random", "tile", and "interleave". '
+ 'See repeat_prompts() in the source code for details.')
+
+ parser.add_argument("--shuffle-seed",
+ type=int,
+ default=0,
+ help='Random seed when the repeat mode is "random"')
+
+ parser = EngineArgs.add_cli_args(parser)
+ args = parser.parse_args()
+ main(args)
diff --git a/benchmarks/benchmark_prefix_caching.py b/benchmarks/benchmark_prefix_caching.py
index 5e9381f712e10..3ab421a89c935 100644
--- a/benchmarks/benchmark_prefix_caching.py
+++ b/benchmarks/benchmark_prefix_caching.py
@@ -10,7 +10,8 @@
--model meta-llama/Llama-2-7b-chat-hf \
--enable-prefix-caching \
--num-prompts 1 \
- --repeat-count 100
+ --repeat-count 100 \
+ --input-length-range 128:256
ShareGPT example usage:
# This command samples 20 prompts with input lengths
diff --git a/benchmarks/benchmark_serving.py b/benchmarks/benchmark_serving.py
index 3256692142c5e..4eb0e1f8ac903 100644
--- a/benchmarks/benchmark_serving.py
+++ b/benchmarks/benchmark_serving.py
@@ -781,6 +781,7 @@ def main(args: argparse.Namespace):
backend = args.backend
model_id = args.model
tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model
+ tokenizer_mode = args.tokenizer_mode
if args.base_url is not None:
api_url = f"{args.base_url}{args.endpoint}"
@@ -790,6 +791,7 @@ def main(args: argparse.Namespace):
base_url = f"http://{args.host}:{args.port}"
tokenizer = get_tokenizer(tokenizer_id,
+ tokenizer_mode=tokenizer_mode,
trust_remote_code=args.trust_remote_code)
if args.dataset is not None:
@@ -1210,5 +1212,15 @@ def main(args: argparse.Namespace):
"from the sampled HF dataset.",
)
+ parser.add_argument(
+ '--tokenizer-mode',
+ type=str,
+ default="auto",
+ choices=['auto', 'slow', 'mistral'],
+ help='The tokenizer mode.\n\n* "auto" will use the '
+ 'fast tokenizer if available.\n* "slow" will '
+ 'always use the slow tokenizer. \n* '
+ '"mistral" will always use the `mistral_common` tokenizer.')
+
args = parser.parse_args()
main(args)
diff --git a/benchmarks/benchmark_serving_guided.py b/benchmarks/benchmark_serving_guided.py
new file mode 100644
index 0000000000000..4435d87e18a8a
--- /dev/null
+++ b/benchmarks/benchmark_serving_guided.py
@@ -0,0 +1,881 @@
+r"""Benchmark online serving throughput with guided decoding.
+
+On the server side, run one of the following commands:
+ (vLLM OpenAI API server)
+ vllm serve --disable-log-requests
+
+ (TGI backend)
+ ./launch_tgi_server.sh
+
+On the client side, run:
+ python benchmarks/benchmark_serving.py \
+ --backend \
+ --model \
+ --dataset json \
+ --guided-decoding-ratio 1.0 \
+ --guided-decoding-backend xgrammar \
+ --request-rate 10 \
+ --num-prompts 1000
+
+ when using tgi backend, add
+ --endpoint /generate_stream
+ to the end of the command above.
+"""
+import argparse
+import asyncio
+import dataclasses
+import json
+import os
+import random
+import time
+import warnings
+from dataclasses import dataclass
+from typing import AsyncGenerator, List, Optional, Tuple
+
+import datasets
+import numpy as np
+import pandas as pd
+from backend_request_func import (ASYNC_REQUEST_FUNCS, RequestFuncInput,
+ RequestFuncOutput)
+from tqdm.asyncio import tqdm
+from transformers import PreTrainedTokenizerBase
+
+try:
+ from vllm.transformers_utils.tokenizer import get_tokenizer
+except ImportError:
+ from backend_request_func import get_tokenizer
+
+try:
+ from vllm.utils import FlexibleArgumentParser
+except ImportError:
+ from argparse import ArgumentParser as FlexibleArgumentParser
+
+MILLISECONDS_TO_SECONDS_CONVERSION = 1000
+
+
+@dataclass
+class BenchmarkMetrics:
+ completed: int
+ total_input: int
+ total_output: int
+ request_throughput: float
+ request_goodput: float
+ output_throughput: float
+ total_token_throughput: float
+ mean_ttft_ms: float
+ median_ttft_ms: float
+ std_ttft_ms: float
+ percentiles_ttft_ms: List[Tuple[float, float]]
+ mean_tpot_ms: float
+ median_tpot_ms: float
+ std_tpot_ms: float
+ percentiles_tpot_ms: List[Tuple[float, float]]
+ mean_itl_ms: float
+ median_itl_ms: float
+ std_itl_ms: float
+ percentiles_itl_ms: List[Tuple[float, float]]
+ # E2EL stands for end-to-end latency per request.
+ # It is the time taken on the client side from sending
+ # a request to receiving a complete response.
+ mean_e2el_ms: float
+ median_e2el_ms: float
+ std_e2el_ms: float
+ percentiles_e2el_ms: List[Tuple[float, float]]
+
+
+@dataclasses.dataclass
+class SampleRequest:
+ """A class representing a single inference request for benchmarking.
+
+ Attributes:
+ prompt: The input text prompt for the model.
+ multi_modal_data: Optional dictionary containing multi-modal data (e.g.
+ images).
+ prompt_len: The length of the prompt in tokens.
+ expected_output_len: The expected length of the output in tokens.
+ """
+ prompt: str
+ prompt_len: int
+ expected_output_len: int
+ schema: dict
+ structure_type: str
+ completion: str = None
+
+
+def sample_requests(tokenizer: PreTrainedTokenizerBase,
+ args: argparse.Namespace) -> List[SampleRequest]:
+ if args.dataset == 'json':
+ if args.json_schema_path is None:
+ dir_path = os.path.dirname(os.path.realpath(__file__))
+ args.json_schema_path = os.path.join(dir_path,
+ "structured_schemas",
+ "structured_schema_1.json")
+ with open(args.json_schema_path) as f:
+ schema = json.load(f)
+ prompt = f"Generate an example of a user profile given the following schema: {json.dumps(schema)}" # noqa: E501
+ input_len = len(tokenizer(prompt).input_ids)
+ print(f"Input length of the prompt: {input_len} tokens")
+ requests = [
+ SampleRequest(prompt=prompt,
+ prompt_len=input_len,
+ expected_output_len=args.output_len,
+ schema=schema,
+ structure_type=args.structure_type)
+ for _ in range(args.num_prompts)
+ ]
+
+ elif args.dataset == "grammar":
+ schema = """
+ ?start: select_statement
+
+ ?select_statement: "SELECT " column_list " FROM " table_name
+
+ ?column_list: column_name ("," column_name)*
+
+ ?table_name: identifier
+
+ ?column_name: identifier
+
+ ?identifier: /[a-zA-Z_][a-zA-Z0-9_]*/
+ """
+ prompt = "Generate an SQL query to show the 'username' \
+ and 'email' from the 'users' table."
+
+ input_len = len(tokenizer(prompt).input_ids)
+ print(f"Input length of the prompt: {input_len} tokens")
+ requests = [
+ SampleRequest(prompt=prompt,
+ prompt_len=input_len,
+ expected_output_len=args.output_len,
+ schema=schema,
+ structure_type=args.structure_type)
+ for _ in range(args.num_prompts)
+ ]
+
+ elif args.dataset == "regex":
+ regex = r"\w+@\w+\.com\n"
+ args.regex = regex
+ prompt = "Generate an email address for Alan Turing, \
+ who works in Enigma. End in .com and new line. \
+ Example result: alan.turing@enigma.com\n"
+
+ input_len = len(tokenizer(prompt).input_ids)
+ print(f"Input length of the prompt: {input_len} tokens")
+ requests = [
+ SampleRequest(prompt=prompt,
+ prompt_len=input_len,
+ expected_output_len=args.output_len,
+ schema=regex,
+ structure_type=args.structure_type)
+ for _ in range(args.num_prompts)
+ ]
+
+ elif args.dataset == "choice":
+ choice = ["Positive", "Negative"]
+ args.choice = choice
+ prompt = "Classify this sentiment: vLLM is wonderful!"
+ input_len = len(tokenizer(prompt).input_ids)
+ print(f"Input length of the prompt: {input_len} tokens")
+ requests = [
+ SampleRequest(prompt=prompt,
+ prompt_len=input_len,
+ expected_output_len=args.output_len,
+ schema=choice,
+ structure_type=args.structure_type)
+ for _ in range(args.num_prompts)
+ ]
+
+ elif args.dataset == "xgrammar_bench":
+ requests: List[SampleRequest] = []
+ dataset = datasets.load_dataset("NousResearch/json-mode-eval",
+ split="train")
+ print(f"dataset has {len(dataset)} entries")
+ len_dataset = len(dataset)
+ for data_point_idx in range(args.num_prompts):
+ idx = data_point_idx
+ while idx >= len_dataset:
+ idx -= len_dataset
+ schema = dataset["schema"][idx]
+ prompt = tokenizer.apply_chat_template(dataset["prompt"][idx],
+ tokenize=False)
+ input_len = len(tokenizer(prompt).input_ids)
+ completion = dataset["completion"][idx]
+
+ requests.append(
+ SampleRequest(prompt=prompt,
+ prompt_len=input_len,
+ expected_output_len=args.output_len,
+ schema=schema,
+ structure_type=args.structure_type,
+ completion=completion))
+
+ return requests
+
+
+async def get_request(
+ input_requests: List[SampleRequest],
+ request_rate: float,
+ burstiness: float = 1.0,
+) -> AsyncGenerator[Tuple[int, SampleRequest], None]:
+ """
+ Asynchronously generates requests at a specified rate
+ with OPTIONAL burstiness.
+
+ Args:
+ input_requests:
+ A list of input requests, each represented as a tuple.
+ request_rate:
+ The rate at which requests are generated (requests/s).
+ burstiness (optional):
+ The burstiness factor of the request generation.
+ Only takes effect when request_rate is not inf.
+ Default value is 1, which follows a Poisson process.
+ Otherwise, the request intervals follow a gamma distribution.
+ A lower burstiness value (0 < burstiness < 1) results
+ in more bursty requests, while a higher burstiness value
+ (burstiness > 1) results in a more uniform arrival of requests.
+ """
+ input_requests = iter(input_requests)
+
+ # Calculate scale parameter theta to maintain the desired request_rate.
+ assert burstiness > 0, (
+ f"A positive burstiness factor is expected, but given {burstiness}.")
+ theta = 1.0 / (request_rate * burstiness)
+
+ for i, request in enumerate(input_requests):
+ yield i, request
+
+ if request_rate == float("inf"):
+ # If the request rate is infinity, then we don't need to wait.
+ continue
+
+ # Sample the request interval from the gamma distribution.
+ # If burstiness is 1, it follows exponential distribution.
+ interval = np.random.gamma(shape=burstiness, scale=theta)
+ # The next request will be sent after the interval.
+ await asyncio.sleep(interval)
+
+
+def calculate_metrics(
+ input_requests: List[Tuple[str, int, int]],
+ outputs: List[RequestFuncOutput],
+ dur_s: float,
+ tokenizer: PreTrainedTokenizerBase,
+ selected_percentile_metrics: List[str],
+ selected_percentiles: List[float],
+) -> Tuple[BenchmarkMetrics, List[int]]:
+ actual_output_lens: List[int] = []
+ total_input = 0
+ completed = 0
+ good_completed = 0
+ itls: List[float] = []
+ tpots: List[float] = []
+ all_tpots: List[float] = []
+ ttfts: List[float] = []
+ e2els: List[float] = []
+ for i in range(len(outputs)):
+ if outputs[i].success:
+ # We use the tokenizer to count the number of output tokens for all
+ # serving backends instead of looking at len(outputs[i].itl) since
+ # multiple output tokens may be bundled together
+ # Note : this may inflate the output token count slightly
+ output_len = len(
+ tokenizer(outputs[i].generated_text,
+ add_special_tokens=False).input_ids)
+ actual_output_lens.append(output_len)
+ total_input += input_requests[i].prompt_len
+ tpot = 0
+ if output_len > 1:
+ tpot = (outputs[i].latency - outputs[i].ttft) / (output_len -
+ 1)
+ tpots.append(tpot)
+ outputs[i].tpot = sum(tpots) / len(tpots) if len(tpots) else 0
+ # Note: if output_len <= 1, we regard tpot as 0 for goodput
+ all_tpots.append(tpot)
+ itls += outputs[i].itl
+ ttfts.append(outputs[i].ttft)
+ e2els.append(outputs[i].latency)
+ completed += 1
+ else:
+ actual_output_lens.append(0)
+
+ if completed == 0:
+ warnings.warn(
+ "All requests failed. This is likely due to a misconfiguration "
+ "on the benchmark arguments.",
+ stacklevel=2)
+ metrics = BenchmarkMetrics(
+ completed=completed,
+ total_input=total_input,
+ total_output=sum(actual_output_lens),
+ request_throughput=completed / dur_s,
+ request_goodput=good_completed / dur_s,
+ output_throughput=sum(actual_output_lens) / dur_s,
+ total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s,
+ mean_ttft_ms=np.mean(ttfts or 0) *
+ 1000, # ttfts is empty if streaming is not supported by backend
+ std_ttft_ms=np.std(ttfts or 0) * 1000,
+ median_ttft_ms=np.median(ttfts or 0) * 1000,
+ percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000)
+ for p in selected_percentiles],
+ mean_tpot_ms=np.mean(tpots or 0) * 1000,
+ std_tpot_ms=np.std(tpots or 0) * 1000,
+ median_tpot_ms=np.median(tpots or 0) * 1000,
+ percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000)
+ for p in selected_percentiles],
+ mean_itl_ms=np.mean(itls or 0) * 1000,
+ std_itl_ms=np.std(itls or 0) * 1000,
+ median_itl_ms=np.median(itls or 0) * 1000,
+ percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000)
+ for p in selected_percentiles],
+ mean_e2el_ms=np.mean(e2els or 0) * 1000,
+ std_e2el_ms=np.std(e2els or 0) * 1000,
+ median_e2el_ms=np.median(e2els or 0) * 1000,
+ percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000)
+ for p in selected_percentiles],
+ )
+
+ return metrics, actual_output_lens
+
+
+async def benchmark(
+ backend: str,
+ api_url: str,
+ base_url: str,
+ model_id: str,
+ tokenizer: PreTrainedTokenizerBase,
+ input_requests: List[SampleRequest],
+ request_rate: float,
+ burstiness: float,
+ disable_tqdm: bool,
+ profile: bool,
+ selected_percentile_metrics: List[str],
+ selected_percentiles: List[str],
+ ignore_eos: bool,
+ max_concurrency: Optional[int],
+ guided_decoding_ratio: float,
+ guided_decoding_backend: str,
+):
+ if backend in ASYNC_REQUEST_FUNCS:
+ request_func = ASYNC_REQUEST_FUNCS[backend]
+ else:
+ raise ValueError(f"Unknown backend: {backend}")
+
+ def prepare_extra_body(request) -> dict:
+ extra_body = {}
+ # Add the schema to the extra_body
+ extra_body[request.structure_type] = request.schema
+ # Add the specific guided_decoding_backend
+ extra_body["guided_decoding_backend"] = guided_decoding_backend
+ return extra_body
+
+ print("Starting initial single prompt test run...")
+ guided_decoding_req_idx = random.sample(
+ range(len(input_requests)),
+ int(len(input_requests) * guided_decoding_ratio))
+
+ test_request = input_requests[0]
+ test_input = RequestFuncInput(
+ model=model_id,
+ prompt=test_request.prompt,
+ api_url=api_url,
+ prompt_len=test_request.prompt_len,
+ output_len=test_request.expected_output_len,
+ ignore_eos=ignore_eos,
+ extra_body=prepare_extra_body(test_request),
+ )
+ test_output = await request_func(request_func_input=test_input)
+ if not test_output.success:
+ raise ValueError(
+ "Initial test run failed - Please make sure benchmark arguments "
+ f"are correctly specified. Error: {test_output.error}")
+ else:
+ print("Initial test run completed. Starting main benchmark run...")
+
+ if profile:
+ print("Starting profiler...")
+ profile_input = RequestFuncInput(
+ model=model_id,
+ prompt=test_request.prompt,
+ api_url=base_url + "/start_profile",
+ prompt_len=test_request.prompt_len,
+ output_len=test_request.expected_output_len,
+ ignore_eos=ignore_eos,
+ extra_body=prepare_extra_body(test_request),
+ )
+ profile_output = await request_func(request_func_input=profile_input)
+ if profile_output.success:
+ print("Profiler started")
+
+ if burstiness == 1.0:
+ distribution = "Poisson process"
+ else:
+ distribution = "Gamma distribution"
+
+ print(f"Traffic request rate: {request_rate}")
+ print(f"Burstiness factor: {burstiness} ({distribution})")
+ print(f"Maximum request concurrency: {max_concurrency}")
+
+ pbar = None if disable_tqdm else tqdm(total=len(input_requests))
+
+ # This can be used once the minimum Python version is 3.10 or higher,
+ # and it will simplify the code in limited_request_func.
+ # semaphore = (asyncio.Semaphore(max_concurrency)
+ # if max_concurrency else contextlib.nullcontext())
+ semaphore = (asyncio.Semaphore(max_concurrency)
+ if max_concurrency else None)
+
+ async def limited_request_func(request_func_input, pbar):
+ if semaphore is None:
+ return await request_func(request_func_input=request_func_input,
+ pbar=pbar)
+ async with semaphore:
+ return await request_func(request_func_input=request_func_input,
+ pbar=pbar)
+
+ benchmark_start_time = time.perf_counter()
+ tasks: List[asyncio.Task] = []
+ expected: List[str] = []
+ async for i, request in get_request(input_requests, request_rate,
+ burstiness):
+ extra_body = prepare_extra_body(
+ request) if i in guided_decoding_req_idx else None
+ request_func_input = RequestFuncInput(
+ model=model_id,
+ prompt=request.prompt,
+ api_url=api_url,
+ prompt_len=request.prompt_len,
+ output_len=request.expected_output_len,
+ ignore_eos=ignore_eos,
+ extra_body=extra_body,
+ )
+ expected.append(request.completion)
+ tasks.append(
+ asyncio.create_task(
+ limited_request_func(request_func_input=request_func_input,
+ pbar=pbar)))
+ outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks)
+
+ if profile:
+ print("Stopping profiler...")
+ profile_input = RequestFuncInput(
+ model=model_id,
+ prompt=test_request.prompt,
+ api_url=base_url + "/stop_profile",
+ prompt_len=test_request.prompt_len,
+ output_len=test_request.expected_output_len,
+ extra_body={test_request.structure_type: test_request.schema},
+ )
+ profile_output = await request_func(request_func_input=profile_input)
+ if profile_output.success:
+ print("Profiler stopped")
+
+ if pbar is not None:
+ pbar.close()
+
+ benchmark_duration = time.perf_counter() - benchmark_start_time
+
+ metrics, actual_output_lens = calculate_metrics(
+ input_requests=input_requests,
+ outputs=outputs,
+ dur_s=benchmark_duration,
+ tokenizer=tokenizer,
+ selected_percentile_metrics=selected_percentile_metrics,
+ selected_percentiles=selected_percentiles,
+ )
+
+ print("{s:{c}^{n}}".format(s=' Serving Benchmark Result ', n=50, c='='))
+ print("{:<40} {:<10}".format("Successful requests:", metrics.completed))
+ print("{:<40} {:<10.2f}".format("Benchmark duration (s):",
+ benchmark_duration))
+ print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input))
+ print("{:<40} {:<10}".format("Total generated tokens:",
+ metrics.total_output))
+ print("{:<40} {:<10.2f}".format("Request throughput (req/s):",
+ metrics.request_throughput))
+ print("{:<40} {:<10.2f}".format("Output token throughput (tok/s):",
+ metrics.output_throughput))
+ print("{:<40} {:<10.2f}".format("Total Token throughput (tok/s):",
+ metrics.total_token_throughput))
+
+ result = {
+ "duration":
+ benchmark_duration,
+ "completed":
+ metrics.completed,
+ "total_input_tokens":
+ metrics.total_input,
+ "total_output_tokens":
+ metrics.total_output,
+ "request_throughput":
+ metrics.request_throughput,
+ "output_throughput":
+ metrics.output_throughput,
+ "total_token_throughput":
+ metrics.total_token_throughput,
+ "ttft_description":
+ pd.Series([output.ttft for output in outputs]).describe().to_dict(),
+ "tpot_description":
+ pd.Series([output.tpot for output in outputs]).describe().to_dict(),
+ "input_lens": [output.prompt_len for output in outputs],
+ "output_lens":
+ actual_output_lens,
+ "ttfts": [output.ttft for output in outputs],
+ "itls": [output.itl for output in outputs],
+ "errors": [output.error for output in outputs],
+ }
+
+ ret = [{
+ 'generated': output.generated_text,
+ 'expected': gt
+ } for output, gt in zip(outputs, expected)]
+
+ def process_one_metric(
+ # E.g., "ttft"
+ metric_attribute_name: str,
+ # E.g., "TTFT"
+ metric_name: str,
+ # E.g., "Time to First Token"
+ metric_header: str,
+ ):
+ # This function prints and adds statistics of the specified
+ # metric.
+ if metric_attribute_name not in selected_percentile_metrics:
+ return
+ print("{s:{c}^{n}}".format(s=metric_header, n=50, c='-'))
+ print("{:<40} {:<10.2f}".format(
+ f"Mean {metric_name} (ms):",
+ getattr(metrics, f"mean_{metric_attribute_name}_ms")))
+ print("{:<40} {:<10.2f}".format(
+ f"Median {metric_name} (ms):",
+ getattr(metrics, f"median_{metric_attribute_name}_ms")))
+ result[f"mean_{metric_attribute_name}_ms"] = getattr(
+ metrics, f"mean_{metric_attribute_name}_ms")
+ result[f"median_{metric_attribute_name}_ms"] = getattr(
+ metrics, f"median_{metric_attribute_name}_ms")
+ result[f"std_{metric_attribute_name}_ms"] = getattr(
+ metrics, f"std_{metric_attribute_name}_ms")
+ for p, value in getattr(metrics,
+ f"percentiles_{metric_attribute_name}_ms"):
+ p_word = str(int(p)) if int(p) == p else str(p)
+ print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):",
+ value))
+ result[f"p{p_word}_{metric_attribute_name}_ms"] = value
+
+ process_one_metric("ttft", "TTFT", "Time to First Token")
+ process_one_metric("tpot", "TPOT",
+ "Time per Output Token (excl. 1st token)")
+ process_one_metric("itl", "ITL", "Inter-token Latency")
+ process_one_metric("e2el", "E2EL", "End-to-end Latency")
+
+ print("=" * 50)
+
+ return result, ret
+
+
+def evaluate(ret, args):
+
+ def _eval_correctness_json(expected, actual):
+ # extract json string from string using regex
+ import re
+ actual = actual.replace('\n', '').replace(' ', '').strip()
+ try:
+ actual = re.search(r'\{.*\}', actual).group()
+ actual = json.loads(actual)
+ except Exception:
+ return False
+
+ return True
+
+ def _eval_correctness_choice(expected, actual):
+ return actual in args.choice
+
+ def _eval_correctness_regex(expected, actual):
+ import re
+ return re.match(args.regex, actual) is not None
+
+ def _eval_correctness(expected, actual):
+ if args.structure_type == 'guided_json':
+ return _eval_correctness_json(expected, actual)
+ elif args.structure_type == 'guided_regex':
+ return _eval_correctness_regex(expected, actual)
+ elif args.structure_type == 'guided_choice':
+ return _eval_correctness_choice(expected, actual)
+ else:
+ return None
+
+ scores = []
+ for res in ret:
+ score = _eval_correctness(res['expected'], res['generated'])
+ res['correctness'] = score
+ scores.append(score)
+
+ not_none_scores = [score for score in scores if score is not None]
+
+ return (sum(not_none_scores) / len(not_none_scores) *
+ 100) if len(not_none_scores) > 0 else None
+
+
+def main(args: argparse.Namespace):
+ print(args)
+ random.seed(args.seed)
+ np.random.seed(args.seed)
+
+ backend = args.backend
+ model_id = args.model
+ tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model
+
+ if args.base_url is not None:
+ api_url = f"{args.base_url}{args.endpoint}"
+ base_url = f"{args.base_url}"
+ else:
+ api_url = f"http://{args.host}:{args.port}{args.endpoint}"
+ base_url = f"http://{args.host}:{args.port}"
+
+ tokenizer = get_tokenizer(tokenizer_id,
+ trust_remote_code=args.trust_remote_code)
+
+ if args.dataset == 'grammar':
+ args.structure_type = 'guided_grammar'
+ elif args.dataset == 'regex':
+ args.structure_type = 'guided_regex'
+ elif args.dataset == 'choice':
+ args.structure_type = 'guided_choice'
+ else:
+ args.structure_type = 'guided_json'
+
+ if args.no_guided_decoding:
+ args.guided_decoding_ratio = 0
+ if args.save_results:
+ result_file_name = f'{args.guided_decoding_ratio}guided'
+ result_file_name += f"_{backend}"
+ result_file_name += f"_{args.request_rate}qps"
+ result_file_name += f"_{args.model.split('/')[-1]}"
+ result_file_name += f"_{args.dataset}"
+ result_file_name += f"_{args.num_prompts}"
+ result_file_name += f"_out{args.output_len}"
+ result_file_name += ".txt"
+ else:
+ result_file_name = None
+
+ input_requests = sample_requests(tokenizer, args)
+
+ benchmark_result, ret = asyncio.run(
+ benchmark(
+ backend=backend,
+ api_url=api_url,
+ base_url=base_url,
+ model_id=model_id,
+ tokenizer=tokenizer,
+ input_requests=input_requests,
+ request_rate=args.request_rate,
+ burstiness=args.burstiness,
+ disable_tqdm=args.disable_tqdm,
+ profile=args.profile,
+ selected_percentile_metrics=args.percentile_metrics.split(","),
+ selected_percentiles=[
+ float(p) for p in args.metric_percentiles.split(",")
+ ],
+ ignore_eos=args.ignore_eos,
+ max_concurrency=args.max_concurrency,
+ guided_decoding_ratio=args.guided_decoding_ratio,
+ guided_decoding_backend=args.guided_decoding_backend,
+ ))
+
+ # Save config and results to json
+ score = evaluate(ret, args)
+ print("correct_rate(%)", score, '\n')
+ if args.save_results:
+ results = {
+ "backend":
+ backend,
+ "model_id":
+ model_id,
+ "tokenizer_id":
+ tokenizer_id,
+ "num_prompts":
+ args.num_prompts,
+ "request_rate":
+ args.request_rate if args.request_rate < float("inf") else "inf",
+ "burstiness":
+ args.burstiness,
+ "max_concurrency":
+ args.max_concurrency,
+ "correct_rate(%)":
+ score
+ }
+ results = {"outputs": ret, **results, **benchmark_result}
+
+ # Save to file
+ if args.result_filename:
+ result_file_name = args.result_filename
+ if args.result_dir:
+ result_file_name = os.path.join(args.result_dir, result_file_name)
+ with open(result_file_name, "w", encoding='utf-8') as outfile:
+ json.dump(results, outfile, indent=4)
+
+
+if __name__ == "__main__":
+ parser = FlexibleArgumentParser(
+ description="Benchmark the online serving throughput.")
+ parser.add_argument(
+ "--backend",
+ type=str,
+ default="vllm",
+ choices=list(ASYNC_REQUEST_FUNCS.keys()),
+ )
+ parser.add_argument(
+ "--base-url",
+ type=str,
+ default=None,
+ help="Server or API base url if not using http host and port.",
+ )
+ parser.add_argument("--host", type=str, default="localhost")
+ parser.add_argument("--port", type=int, default=8000)
+ parser.add_argument(
+ "--endpoint",
+ type=str,
+ default="/v1/completions",
+ help="API endpoint.",
+ )
+ parser.add_argument(
+ "--dataset",
+ default='json',
+ choices=['json', 'grammar', 'regex', 'choice', 'xgrammar_bench'])
+ parser.add_argument("--json_schema_path",
+ type=str,
+ default=None,
+ help="Path to json schema.")
+ parser.add_argument(
+ "--max-concurrency",
+ type=int,
+ default=None,
+ help="Maximum number of concurrent requests. This can be used "
+ "to help simulate an environment where a higher level component "
+ "is enforcing a maximum number of concurrent requests. While the "
+ "--request-rate argument controls the rate at which requests are "
+ "initiated, this argument will control how many are actually allowed "
+ "to execute at a time. This means that when used in combination, the "
+ "actual request rate may be lower than specified with --request-rate, "
+ "if the server is not processing requests fast enough to keep up.")
+ parser.add_argument(
+ "--model",
+ type=str,
+ required=True,
+ help="Name of the model.",
+ )
+ parser.add_argument(
+ "--tokenizer",
+ type=str,
+ help=
+ "Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501
+ )
+ parser.add_argument(
+ "--num-prompts",
+ type=int,
+ default=1000,
+ help="Number of prompts to process.",
+ )
+ parser.add_argument(
+ "--output-len",
+ type=int,
+ default=128,
+ help="Number of output tokens.",
+ )
+ parser.add_argument(
+ "--request-rate",
+ type=float,
+ default=float("inf"),
+ help="Number of requests per second. If this is inf, "
+ "then all the requests are sent at time 0. "
+ "Otherwise, we use Poisson process or gamma distribution "
+ "to synthesize the request arrival times.",
+ )
+ parser.add_argument(
+ "--burstiness",
+ type=float,
+ default=1.0,
+ help="Burstiness factor of the request generation. "
+ "Only take effect when request_rate is not inf. "
+ "Default value is 1, which follows Poisson process. "
+ "Otherwise, the request intervals follow a gamma distribution. "
+ "A lower burstiness value (0 < burstiness < 1) results in more "
+ "bursty requests. A higher burstiness value (burstiness > 1) "
+ "results in a more uniform arrival of requests.",
+ )
+ parser.add_argument("--seed", type=int, default=0)
+ parser.add_argument(
+ "--trust-remote-code",
+ action="store_true",
+ help="Trust remote code from huggingface",
+ )
+ parser.add_argument(
+ "--disable-tqdm",
+ action="store_true",
+ help="Specify to disable tqdm progress bar.",
+ )
+ parser.add_argument(
+ "--save-results",
+ action="store_true",
+ help="Specify to save benchmark results to a json file",
+ )
+ parser.add_argument(
+ "--profile",
+ action="store_true",
+ help="Use Torch Profiler. The endpoint must be launched with "
+ "VLLM_TORCH_PROFILER_DIR to enable profiler.",
+ )
+ parser.add_argument(
+ "--result-dir",
+ type=str,
+ default=None,
+ help="Specify directory to save benchmark json results."
+ "If not specified, results are saved in the current directory.",
+ )
+ parser.add_argument(
+ "--result-filename",
+ type=str,
+ default=None,
+ help="Specify the filename to save benchmark json results."
+ "If not specified, results will be saved in "
+ "{backend}-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
+ " format.",
+ )
+ parser.add_argument(
+ "--ignore-eos",
+ action="store_true",
+ help="Set ignore_eos flag when sending the benchmark request."
+ "Warning: ignore_eos is not supported in deepspeed_mii and tgi.")
+ parser.add_argument(
+ "--percentile-metrics",
+ type=str,
+ default="ttft,tpot,itl",
+ help="Comma-seperated list of selected metrics to report percentils. "
+ "This argument specifies the metrics to report percentiles. "
+ "Allowed metric names are \"ttft\", \"tpot\", \"itl\", \"e2el\". "
+ "Default value is \"ttft,tpot,itl\".")
+ parser.add_argument(
+ "--metric-percentiles",
+ type=str,
+ default="99",
+ help="Comma-seperated list of percentiles for selected metrics. "
+ "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "
+ "Default value is \"99\". "
+ "Use \"--percentile-metrics\" to select metrics.",
+ )
+ parser.add_argument("--no-guided-decoding",
+ action='store_true',
+ default=False,
+ help="Whether to disable JSON decoding or not.")
+ parser.add_argument("--guided-decoding-ratio",
+ type=float,
+ default=1.0,
+ help="Ratio of Guided Decoding requests")
+ parser.add_argument("--guided-decoding-backend",
+ type=str,
+ choices=["outlines", "lm-format-enforcer", "xgrammar"],
+ default="xgrammar",
+ help="Backend to use for guided decoding")
+
+ args = parser.parse_args()
+ main(args)
diff --git a/benchmarks/benchmark_throughput.py b/benchmarks/benchmark_throughput.py
index 159cf055737ce..c1b10b3cf8f58 100644
--- a/benchmarks/benchmark_throughput.py
+++ b/benchmarks/benchmark_throughput.py
@@ -4,7 +4,8 @@
import json
import random
import time
-from typing import List, Optional
+from functools import cache
+from typing import Dict, List, Optional, Tuple
import torch
import uvloop
@@ -17,8 +18,11 @@
from vllm.entrypoints.openai.api_server import (
build_async_engine_client_from_engine_args)
from vllm.inputs import TextPrompt
+from vllm.lora.request import LoRARequest
+from vllm.lora.utils import get_adapter_absolute_path
from vllm.multimodal import MultiModalDataDict
from vllm.sampling_params import BeamSearchParams
+from vllm.transformers_utils.tokenizer import AnyTokenizer, get_lora_tokenizer
from vllm.utils import FlexibleArgumentParser, merge_async_iterators
@@ -28,15 +32,17 @@ class SampleRequest:
Attributes:
prompt: The input text prompt for the model.
- multi_modal_data: Optional dictionary containing multi-modal data (e.g.
- images).
prompt_len: The length of the prompt in tokens.
expected_output_len: The expected length of the output in tokens.
+ multi_modal_data: Optional dictionary containing multi-modal data (e.g.
+ images).
+ lora_request: Optional LoRARequest specifying the LoRA to use.
"""
prompt: str
prompt_len: int
expected_output_len: int
multi_modal_data: Optional[MultiModalDataDict] = None
+ lora_request: Optional[LoRARequest] = None
def _get_prompt_for_image_model(question: str, *, model: str) -> str:
@@ -60,8 +66,30 @@ def _get_prompt_for_image_model(question: str, *, model: str) -> str:
raise ValueError(f"Unsupported model {model}")
+@cache
+def lora_path_on_disk(lora_path: str) -> str:
+ return get_adapter_absolute_path(lora_path)
+
+
+lora_tokenizer_cache: Dict[int, AnyTokenizer] = {}
+
+
+def get_random_lora_request(
+ args: argparse.Namespace
+) -> Tuple[LoRARequest, Optional[AnyTokenizer]]:
+ global lora_tokenizer_cache
+ lora_id = random.randint(1, args.max_loras)
+ lora_request = LoRARequest(lora_name=str(lora_id),
+ lora_int_id=lora_id,
+ lora_path=lora_path_on_disk(args.lora_path))
+ if lora_id not in lora_tokenizer_cache:
+ lora_tokenizer_cache[lora_id] = get_lora_tokenizer(lora_request)
+ return lora_request, lora_tokenizer_cache[lora_id]
+
+
def sample_requests(tokenizer: PreTrainedTokenizerBase,
args: argparse.Namespace) -> List[SampleRequest]:
+
dataset_path: str = args.dataset
num_requests: int = args.num_prompts
fixed_output_len: Optional[int] = args.output_len
@@ -79,7 +107,9 @@ def sample_requests(tokenizer: PreTrainedTokenizerBase,
# Filter out sequences that are too long or too short
filtered_dataset: List[SampleRequest] = []
- for data in dataset:
+ for data in tqdm(dataset,
+ total=len(filtered_dataset),
+ desc="sampling requests"):
if len(filtered_dataset) == num_requests:
break
@@ -102,9 +132,16 @@ def sample_requests(tokenizer: PreTrainedTokenizerBase,
continue
prompt = _get_prompt_for_image_model(question=prompt, model=model)
+ request_tokenizer = tokenizer
+ lora_request: Optional[LoRARequest] = None
+ if args.enable_lora:
+ lora_request, lora_tokenizer = get_random_lora_request(args)
+ if lora_tokenizer:
+ request_tokenizer = lora_tokenizer
+
# Tokenize the prompts and completions.
- prompt_token_ids = tokenizer(prompt).input_ids
- completion_token_ids = tokenizer(completion).input_ids
+ prompt_token_ids = request_tokenizer(prompt).input_ids
+ completion_token_ids = request_tokenizer(completion).input_ids
prompt_len = len(prompt_token_ids)
output_len = len(completion_token_ids
) if fixed_output_len is None else fixed_output_len
@@ -118,7 +155,8 @@ def sample_requests(tokenizer: PreTrainedTokenizerBase,
SampleRequest(prompt=prompt,
prompt_len=prompt_len,
expected_output_len=output_len,
- multi_modal_data=multi_modal_data))
+ multi_modal_data=multi_modal_data,
+ lora_request=lora_request))
return filtered_dataset
@@ -146,14 +184,21 @@ def run_vllm(
ignore_eos=True,
max_tokens=request.expected_output_len,
))
+ lora_requests: Optional[List[LoRARequest]] = None
+ if engine_args.enable_lora:
+ lora_requests = [request.lora_request for request in requests]
use_beam_search = False
if not use_beam_search:
start = time.perf_counter()
- llm.generate(prompts, sampling_params, use_tqdm=True)
+ llm.generate(prompts,
+ sampling_params,
+ lora_request=lora_requests,
+ use_tqdm=True)
end = time.perf_counter()
else:
+ assert lora_requests is None, "BeamSearch API does not support LoRA"
prompts = [request.prompt for request in requests]
# output_len should be the same for all requests.
output_len = requests[0][2]
@@ -185,6 +230,7 @@ async def run_vllm_async(
# Add the requests to the engine.
prompts: List[TextPrompt] = []
sampling_params: List[SamplingParams] = []
+ lora_requests: List[Optional[LoRARequest]] = []
for request in requests:
prompts.append(
TextPrompt(prompt=request.prompt,
@@ -197,11 +243,16 @@ async def run_vllm_async(
ignore_eos=True,
max_tokens=request.expected_output_len,
))
+ lora_requests.append(request.lora_request)
generators = []
start = time.perf_counter()
- for i, (prompt, sp) in enumerate(zip(prompts, sampling_params)):
- generator = llm.generate(prompt, sp, request_id=f"test{i}")
+ for i, (prompt, sp,
+ lr) in enumerate(zip(prompts, sampling_params, lora_requests)):
+ generator = llm.generate(prompt,
+ sp,
+ lora_request=lr,
+ request_id=f"test{i}")
generators.append(generator)
all_gens = merge_async_iterators(*generators)
async for i, res in all_gens:
@@ -294,23 +345,45 @@ def main(args: argparse.Namespace):
tokenizer = AutoTokenizer.from_pretrained(
args.tokenizer, trust_remote_code=args.trust_remote_code)
if args.dataset is None:
- # Synthesize a prompt with the given input length.
- # As tokenizer may add additional tokens like BOS, we need to try
- # different lengths to get the desired input length.
- for i in range(-10, 10):
- prompt = "hi " * (args.input_len + i)
- tokenized_prompt = tokenizer(prompt).input_ids
- if len(tokenized_prompt) == args.input_len:
- break
- else:
- raise ValueError(
- f"Failed to synthesize a prompt with {args.input_len} tokens.")
- requests = [
- SampleRequest(prompt=prompt,
- prompt_len=args.input_len,
- expected_output_len=args.output_len)
- for _ in range(args.num_prompts)
- ]
+ vocab_size = tokenizer.vocab_size
+ requests = []
+ for _ in range(args.num_prompts):
+
+ request_tokenizer = tokenizer
+ lora_request: Optional[LoRARequest] = None
+ if args.enable_lora:
+ lora_request, lora_tokenizer = get_random_lora_request(args)
+ if lora_tokenizer:
+ request_tokenizer = lora_tokenizer
+
+ # Synthesize a prompt with the given input length.
+ candidate_ids = [
+ random.randint(0, vocab_size - 1)
+ for _ in range(args.input_len)
+ ]
+ # As tokenizer may add additional tokens like BOS, we need to try
+ # different lengths to get the desired input length.
+ for _ in range(5): # Max attempts to correct
+ candidate_prompt = request_tokenizer.decode(candidate_ids)
+ tokenized_len = len(request_tokenizer.encode(candidate_prompt))
+
+ if tokenized_len == args.input_len:
+ break
+
+ # Adjust length based on difference
+ diff = args.input_len - tokenized_len
+ if diff > 0:
+ candidate_ids.extend([
+ random.randint(100, vocab_size - 100)
+ for _ in range(diff)
+ ])
+ else:
+ candidate_ids = candidate_ids[:diff]
+ requests.append(
+ SampleRequest(prompt=candidate_prompt,
+ prompt_len=args.input_len,
+ expected_output_len=args.output_len,
+ lora_request=lora_request))
else:
requests = sample_requests(tokenizer, args)
@@ -409,6 +482,14 @@ def main(args: argparse.Namespace):
action='store_true',
default=False,
help="Disable decoupled async engine frontend.")
+ # LoRA
+ parser.add_argument(
+ "--lora-path",
+ type=str,
+ default=None,
+ help="Path to the lora adapters to use. This can be an absolute path, "
+ "a relative path, or a Hugging Face model identifier.")
+
parser = AsyncEngineArgs.add_cli_args(parser)
args = parser.parse_args()
if args.tokenizer is None:
@@ -418,6 +499,8 @@ def main(args: argparse.Namespace):
assert args.output_len is not None
else:
assert args.input_len is None
+ if args.enable_lora:
+ assert args.lora_path is not None
if args.backend == "vllm":
if args.hf_max_batch_size is not None:
@@ -427,6 +510,9 @@ def main(args: argparse.Namespace):
raise ValueError("HF max batch size is required for HF backend.")
if args.quantization is not None:
raise ValueError("Quantization is only for vLLM backend.")
+ if args.enable_lora is not None:
+ raise ValueError("LoRA benchmarking is only supported for vLLM"
+ " backend")
elif args.backend == "mii":
if args.dtype != "auto":
raise ValueError("dtype must be auto for MII backend.")
@@ -439,4 +525,7 @@ def main(args: argparse.Namespace):
if args.tokenizer != args.model:
raise ValueError("Tokenizer must be the same as the model for MII "
"backend.")
+ if args.enable_lora is not None:
+ raise ValueError("LoRA benchmarking is only supported for vLLM"
+ " backend")
main(args)
diff --git a/benchmarks/cutlass_benchmarks/sparse_benchmarks.py b/benchmarks/cutlass_benchmarks/sparse_benchmarks.py
new file mode 100644
index 0000000000000..3d1c5e392f9e2
--- /dev/null
+++ b/benchmarks/cutlass_benchmarks/sparse_benchmarks.py
@@ -0,0 +1,384 @@
+import argparse
+import copy
+import itertools
+import pickle as pkl
+import time
+from typing import Callable, Iterable, List, Tuple
+
+import torch
+import torch.utils.benchmark as TBenchmark
+from torch.utils.benchmark import Measurement as TMeasurement
+from utils import make_rand_sparse_tensors
+from weight_shapes import WEIGHT_SHAPES
+
+from vllm import _custom_ops as ops
+from vllm.utils import FlexibleArgumentParser
+
+DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())
+DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512]
+DEFAULT_TP_SIZES = [1]
+
+
+# bench
+def bench_fn(label: str, sub_label: str, description: str, fn: Callable, *args,
+ **kwargs) -> TMeasurement:
+ min_run_time = 1
+
+ globals = {
+ "args": args,
+ "kwargs": kwargs,
+ "fn": fn,
+ }
+ return TBenchmark.Timer(
+ stmt="fn(*args, **kwargs)",
+ globals=globals,
+ label=label,
+ sub_label=sub_label,
+ description=description,
+ ).blocked_autorange(min_run_time=min_run_time)
+
+
+def bench_int8(dtype: torch.dtype, m: int, k: int, n: int, label: str,
+ sub_label: str) -> Iterable[TMeasurement]:
+ assert dtype == torch.int8
+ b_compressed, e, a, b = make_rand_sparse_tensors(torch.int8, m, n, k)
+ scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
+ scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
+ bias = torch.zeros((n, ), device="cuda", dtype=torch.bfloat16)
+
+ out = ops.cutlass_scaled_sparse_mm(a, b_compressed, e, scale_a, scale_b,
+ torch.bfloat16)
+ out_ref = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)
+
+ if not torch.allclose(out, out_ref):
+ print("Incorrect results")
+ print(out)
+ print(out_ref)
+ else:
+ print("Correct results")
+
+ timers = []
+ # pytorch impl - bfloat16
+ timers.append(
+ bench_fn(label, sub_label, "pytorch_bf16_bf16_bf16_matmul-no-scales",
+ torch.mm, a.to(dtype=torch.bfloat16),
+ b.to(dtype=torch.bfloat16)))
+
+ # pytorch impl - float16
+ timers.append(
+ bench_fn(label, sub_label,
+ "pytorch_fp16_fp16_fp16_matmul-no-scales", torch.mm,
+ a.to(dtype=torch.float16), b.to(dtype=torch.float16)))
+
+ # cutlass impl
+ timers.append(
+ bench_fn(label, sub_label, "cutlass_i8_i8_bf16_scaled_mm",
+ ops.cutlass_scaled_mm, a, b, scale_a, scale_b,
+ torch.bfloat16))
+
+ # cutlass with bias
+ timers.append(
+ bench_fn(label, sub_label, "cutlass_i8_i8_bf16_scaled_mm_bias",
+ ops.cutlass_scaled_mm, a, b, scale_a, scale_b, torch.bfloat16,
+ bias))
+
+ # cutlass sparse impl
+ timers.append(
+ bench_fn(label, sub_label, "cutlass_i8_i8_bf16_scaled_sparse_mm",
+ ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,
+ scale_b, torch.bfloat16))
+
+ # cutlass sparse with bias
+ timers.append(
+ bench_fn(label, sub_label, "cutlass_i8_i8_bf16_scaled_sparse_mm_bias",
+ ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,
+ scale_b, torch.bfloat16, bias))
+
+ return timers
+
+
+def bench_fp8(dtype: torch.dtype, m: int, k: int, n: int, label: str,
+ sub_label: str) -> Iterable[TMeasurement]:
+ assert dtype == torch.float8_e4m3fn
+ b_compressed, e, a, b = make_rand_sparse_tensors(torch.float8_e4m3fn, m, n,
+ k)
+ scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
+ scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
+ bias = torch.zeros((n, ), device="cuda", dtype=torch.bfloat16)
+
+ out = ops.cutlass_scaled_sparse_mm(a, b_compressed, e, scale_a, scale_b,
+ torch.bfloat16)
+ out_ref = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)
+
+ if not torch.allclose(out, out_ref):
+ print("Incorrect results")
+ print(out)
+ print(out_ref)
+ else:
+ print("Correct results")
+
+ timers = []
+
+ # pytorch impl w. bf16
+ timers.append(
+ bench_fn(label, sub_label, "pytorch_bf16_bf16_bf16_matmul-no-scales",
+ torch.mm, a.to(dtype=torch.bfloat16, device="cuda"),
+ b.to(dtype=torch.bfloat16, device="cuda")))
+
+ # pytorch impl: bf16 output, without fp8 fast accum
+ timers.append(
+ bench_fn(label,
+ sub_label,
+ "pytorch_fp8_fp8_bf16_scaled_mm",
+ torch._scaled_mm,
+ a,
+ b,
+ scale_a=scale_a,
+ scale_b=scale_b,
+ out_dtype=torch.bfloat16))
+
+ # pytorch impl: bf16 output, with fp8 fast accum
+ timers.append(
+ bench_fn(label,
+ sub_label,
+ "pytorch_fp8_fp8_bf16_scaled_mm_fast_accum",
+ torch._scaled_mm,
+ a,
+ b,
+ scale_a=scale_a,
+ scale_b=scale_b,
+ out_dtype=torch.bfloat16,
+ use_fast_accum=True))
+
+ # pytorch impl: fp16 output, without fp8 fast accum
+ timers.append(
+ bench_fn(label,
+ sub_label,
+ "pytorch_fp8_fp8_fp16_scaled_mm",
+ torch._scaled_mm,
+ a,
+ b,
+ scale_a=scale_a,
+ scale_b=scale_b,
+ out_dtype=torch.float16))
+
+ # pytorch impl: fp16 output, with fp8 fast accum
+ timers.append(
+ bench_fn(label,
+ sub_label,
+ "pytorch_fp8_fp8_fp16_scaled_mm_fast_accum",
+ torch._scaled_mm,
+ a,
+ b,
+ scale_a=scale_a,
+ scale_b=scale_b,
+ out_dtype=torch.float16,
+ use_fast_accum=True))
+
+ # cutlass impl: bf16 output
+ timers.append(
+ bench_fn(label, sub_label, "cutlass_fp8_fp8_bf16_scaled_mm",
+ ops.cutlass_scaled_mm, a, b, scale_a, scale_b,
+ torch.bfloat16))
+
+ # cutlass impl: bf16 output
+ timers.append(
+ bench_fn(label, sub_label, "cutlass_fp8_fp8_bf16_scaled_sparse_mm",
+ ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,
+ scale_b, torch.bfloat16))
+
+ # cutlass impl: fp16 output
+ timers.append(
+ bench_fn(label, sub_label, "cutlass_fp8_fp8_fp16_scaled_sparse_mm",
+ ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,
+ scale_b, torch.float16))
+
+ # cutlass impl: bf16 output, with bias
+ timers.append(
+ bench_fn(label, sub_label,
+ "cutlass_fp8_fp8_bf16_scaled_sparse_mm_bias",
+ ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,
+ scale_b, torch.bfloat16, bias))
+
+ # cutlass impl: fp16 output, with bias
+ timers.append(
+ bench_fn(label, sub_label,
+ "cutlass_fp8_fp8_fp16_scaled_sparse_mm_bias",
+ ops.cutlass_scaled_sparse_mm, a, b_compressed, e, scale_a,
+ scale_b, torch.float16, bias.to(dtype=torch.float16)))
+
+ return timers
+
+
+def bench(dtype: torch.dtype, m: int, k: int, n: int, label: str,
+ sub_label: str) -> Iterable[TMeasurement]:
+ if dtype == torch.int8:
+ return bench_int8(dtype, m, k, n, label, sub_label)
+ if dtype == torch.float8_e4m3fn:
+ return bench_fp8(dtype, m, k, n, label, sub_label)
+ raise ValueError("unsupported type")
+
+
+# runner
+def print_timers(timers: Iterable[TMeasurement]):
+ compare = TBenchmark.Compare(timers)
+ compare.print()
+
+
+def run(dtype: torch.dtype,
+ MKNs: Iterable[Tuple[int, int, int]]) -> Iterable[TMeasurement]:
+ results = []
+ for m, k, n in MKNs:
+ timers = bench(dtype, m, k, n, f"scaled-{dtype}-gemm",
+ f"MKN=({m}x{k}x{n})")
+ print_timers(timers)
+ results.extend(timers)
+
+ return results
+
+
+# output makers
+def make_output(data: Iterable[TMeasurement],
+ MKNs: Iterable[Tuple[int, int, int]],
+ base_description: str,
+ timestamp=None):
+ print(f"== All Results {base_description} ====")
+ print_timers(data)
+
+ # pickle all the results
+ timestamp = int(time.time()) if timestamp is None else timestamp
+ with open(f"{base_description}-{timestamp}.pkl", "wb") as f:
+ pkl.dump(data, f)
+
+
+# argparse runners
+
+
+def run_square_bench(args):
+ dim_sizes = list(
+ range(args.dim_start, args.dim_end + 1, args.dim_increment))
+ MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes))
+ data = run(args.dtype, MKNs)
+
+ make_output(data, MKNs, f"square_bench-{args.dtype}")
+
+
+def run_range_bench(args):
+ dim_sizes = list(range(args.dim_start, args.dim_end, args.dim_increment))
+ n = len(dim_sizes)
+ Ms = [args.m_constant] * n if args.m_constant is not None else dim_sizes
+ Ks = [args.k_constant] * n if args.k_constant is not None else dim_sizes
+ Ns = [args.n_constant] * n if args.n_constant is not None else dim_sizes
+ MKNs = list(zip(Ms, Ks, Ns))
+ data = run(args.dtype, MKNs)
+
+ make_output(data, MKNs, f"range_bench-{args.dtype}")
+
+
+def run_model_bench(args):
+ print("Benchmarking models:")
+ for i, model in enumerate(args.models):
+ print(f"[{i}] {model}")
+
+ def model_shapes(model_name: str, tp_size: int) -> List[Tuple[int, int]]:
+ KNs = []
+ for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model_name]):
+ KN[tp_split_dim] = KN[tp_split_dim] // tp_size
+ KNs.append(KN)
+ return KNs
+
+ model_bench_data = []
+ models_tps = list(itertools.product(args.models, args.tp_sizes))
+ for model, tp_size in models_tps:
+ Ms = args.batch_sizes
+ KNs = model_shapes(model, tp_size)
+ MKNs = []
+ for m in Ms:
+ for k, n in KNs:
+ MKNs.append((m, k, n))
+
+ data = run(args.dtype, MKNs)
+ model_bench_data.append(data)
+
+ # Print all results
+ for data, model_tp in zip(model_bench_data, models_tps):
+ model, tp_size = model_tp
+ print(f"== Results {args.dtype} {model}-TP{tp_size} ====")
+ print_timers(data)
+
+ timestamp = int(time.time())
+
+ all_data = []
+ for d in model_bench_data:
+ all_data.extend(d)
+ # pickle all data
+ with open(f"model_bench-{args.dtype}-{timestamp}.pkl", "wb") as f:
+ pkl.dump(all_data, f)
+
+
+if __name__ == '__main__':
+
+ def to_torch_dtype(dt):
+ if dt == "int8":
+ return torch.int8
+ if dt == "fp8":
+ return torch.float8_e4m3fn
+ raise ValueError("unsupported dtype")
+
+ parser = FlexibleArgumentParser(
+ description="""
+Benchmark Cutlass GEMM.
+
+ To run square GEMMs:
+ python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 square_bench --dim-start 128 --dim-end 512 --dim-increment 64
+
+ To run constant N and K and sweep M:
+ python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 range_bench --dim-start 128 --dim-end 512 --dim-increment 64 --n-constant 16384 --k-constant 16384
+
+ To run dimensions from a model:
+ python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 model_bench --models meta-llama/Llama-2-7b-hf --batch-sizes 16 --tp-sizes 1
+
+ Output:
+ - a .pkl file, that is a list of raw torch.benchmark.utils.Measurements for the pytorch and cutlass implementations for the various GEMMs.
+ """, # noqa: E501
+ formatter_class=argparse.RawTextHelpFormatter)
+
+ parser.add_argument("--dtype",
+ type=to_torch_dtype,
+ required=True,
+ help="Available options are ['int8', 'fp8']")
+ subparsers = parser.add_subparsers(dest="cmd")
+
+ square_parser = subparsers.add_parser("square_bench")
+ square_parser.add_argument("--dim-start", type=int, required=True)
+ square_parser.add_argument("--dim-end", type=int, required=True)
+ square_parser.add_argument("--dim-increment", type=int, required=True)
+ square_parser.set_defaults(func=run_square_bench)
+
+ range_parser = subparsers.add_parser("range_bench")
+ range_parser.add_argument("--dim-start", type=int, required=True)
+ range_parser.add_argument("--dim-end", type=int, required=True)
+ range_parser.add_argument("--dim-increment", type=int, required=True)
+ range_parser.add_argument("--m-constant", type=int, default=None)
+ range_parser.add_argument("--n-constant", type=int, default=None)
+ range_parser.add_argument("--k-constant", type=int, default=None)
+ range_parser.set_defaults(func=run_range_bench)
+
+ model_parser = subparsers.add_parser("model_bench")
+ model_parser.add_argument("--models",
+ nargs="+",
+ type=str,
+ default=DEFAULT_MODELS,
+ choices=WEIGHT_SHAPES.keys())
+ model_parser.add_argument("--tp-sizes",
+ nargs="+",
+ type=int,
+ default=DEFAULT_TP_SIZES)
+ model_parser.add_argument("--batch-sizes",
+ nargs="+",
+ type=int,
+ default=DEFAULT_BATCH_SIZES)
+ model_parser.set_defaults(func=run_model_bench)
+
+ args = parser.parse_args()
+ args.func(args)
diff --git a/benchmarks/cutlass_benchmarks/utils.py b/benchmarks/cutlass_benchmarks/utils.py
new file mode 100644
index 0000000000000..ef06fcd6604dd
--- /dev/null
+++ b/benchmarks/cutlass_benchmarks/utils.py
@@ -0,0 +1,96 @@
+# Cutlass bench utils
+from typing import Iterable, Tuple
+
+import torch
+
+import vllm._custom_ops as ops
+
+
+def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
+ finfo = torch.finfo(torch.float8_e4m3fn)
+ return torch.round(tensor.clamp(
+ min=finfo.min, max=finfo.max)).to(dtype=torch.float8_e4m3fn)
+
+
+def to_int8(tensor: torch.Tensor) -> torch.Tensor:
+ return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)
+
+
+def to_bf16(tensor: torch.Tensor) -> torch.Tensor:
+ return tensor.to(dtype=torch.bfloat16)
+
+
+def to_fp16(tensor: torch.Tensor) -> torch.Tensor:
+ return tensor.to(dtype=torch.float16)
+
+
+def make_rand_tensors(dtype: torch.dtype, m: int, n: int,
+ k: int) -> Tuple[torch.Tensor, torch.Tensor]:
+ a = torch.randn((m, k), device='cuda') * 5
+ b = torch.randn((n, k), device='cuda').t() * 5
+
+ if dtype == torch.int8:
+ return to_int8(a), to_int8(b)
+ if dtype == torch.float8_e4m3fn:
+ return to_fp8(a), to_fp8(b)
+
+ raise ValueError("unsupported dtype")
+
+
+def prune_to_2_4(tensor):
+ # Reshape tensor to [N, 4] where N is number of groups of 4
+ original_shape = tensor.shape
+ reshaped = tensor.reshape(-1, 4)
+
+ # Get indices of top 2 absolute values in each group of 4
+ _, indices = torch.topk(torch.abs(reshaped), k=2, dim=1)
+
+ # Create binary mask
+ mask = torch.zeros_like(reshaped)
+ mask.scatter_(dim=1,
+ index=indices,
+ src=torch.ones_like(indices, dtype=mask.dtype))
+
+ # Apply mask and reshape back
+ pruned = reshaped * mask
+
+ # Turn all -0.0 to 0.0
+ pruned[pruned == -0.0] = 0.0
+
+ return pruned.reshape(original_shape)
+
+
+def make_rand_sparse_tensors(dtype: torch.dtype, m: int, n: int,
+ k: int) -> Tuple[torch.Tensor, torch.Tensor]:
+ a = torch.randn((m, k), device='cuda') * 5
+ b = torch.randn((n, k), device='cuda').t() * 5
+
+ b = prune_to_2_4(b.t()).t()
+
+ if dtype == torch.int8:
+ a, b = to_int8(a), to_int8(b)
+ elif dtype == torch.float8_e4m3fn:
+ a, b = to_fp8(a), to_fp8(b)
+ elif dtype == torch.float16:
+ a, b = to_fp16(a), to_fp16(b)
+ elif dtype == torch.bfloat16:
+ a, b = to_bf16(a), to_bf16(b)
+ else:
+ raise ValueError("unsupported dtype")
+
+ b_compressed, e = ops.cutlass_sparse_compress(b.t())
+
+ # Compressed B, Metadata, Original A, B
+ return b_compressed, e, a, b
+
+
+def make_n_rand_sparse_tensors(num_tensors: int, dtype: torch.dtype,
+ m: int, n: int, k: int) -> \
+ Tuple[Iterable[torch.Tensor], Iterable[torch.Tensor]]:
+ ABs = []
+ for _ in range(num_tensors):
+ b_comp, e, a, b = make_rand_sparse_tensors(dtype, m, n, k)
+ if b_comp is not None:
+ ABs.append(make_rand_sparse_tensors(dtype, m, n, k))
+ BComps, Es, As, Bs = zip(*ABs)
+ return list(BComps), list(Es), list(As), list(Bs)
diff --git a/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py b/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py
index 63cf5d50cac75..d0353bc8cb42a 100644
--- a/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py
+++ b/benchmarks/cutlass_benchmarks/w8a8_benchmarks.py
@@ -8,6 +8,7 @@
import torch
import torch.utils.benchmark as TBenchmark
from torch.utils.benchmark import Measurement as TMeasurement
+from utils import make_rand_tensors
from weight_shapes import WEIGHT_SHAPES
from vllm import _custom_ops as ops
@@ -17,31 +18,6 @@
DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512]
DEFAULT_TP_SIZES = [1]
-# helpers
-
-
-def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
- finfo = torch.finfo(torch.float8_e4m3fn)
- return torch.round(tensor.clamp(
- min=finfo.min, max=finfo.max)).to(dtype=torch.float8_e4m3fn)
-
-
-def to_int8(tensor: torch.Tensor) -> torch.Tensor:
- return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)
-
-
-def make_rand_tensors(dtype: torch.dtype, m: int, n: int,
- k: int) -> Tuple[torch.Tensor, torch.Tensor]:
- a = torch.randn((m, k), device='cuda') * 5
- b = torch.randn((n, k), device='cuda').t() * 5
-
- if dtype == torch.int8:
- return to_int8(a), to_int8(b)
- if dtype == torch.float8_e4m3fn:
- return to_fp8(a), to_fp8(b)
-
- raise ValueError("unsupported dtype")
-
# bench
def bench_fn(label: str, sub_label: str, description: str, fn: Callable, *args,
@@ -386,4 +362,4 @@ def to_torch_dtype(dt):
model_parser.set_defaults(func=run_model_bench)
args = parser.parse_args()
- args.func(args)
+ args.func(args)
\ No newline at end of file
diff --git a/benchmarks/cutlass_benchmarks/weight_shapes.py b/benchmarks/cutlass_benchmarks/weight_shapes.py
index 25ec9d6028627..d58fb0bf86374 100644
--- a/benchmarks/cutlass_benchmarks/weight_shapes.py
+++ b/benchmarks/cutlass_benchmarks/weight_shapes.py
@@ -40,4 +40,4 @@
([8192, 57344], 1),
([28672, 8192], 0),
],
-}
+}
\ No newline at end of file
diff --git a/benchmarks/disagg_benchmarks/disagg_overhead_benchmark.sh b/benchmarks/disagg_benchmarks/disagg_overhead_benchmark.sh
index 2924ea4a49f54..94999630bae12 100644
--- a/benchmarks/disagg_benchmarks/disagg_overhead_benchmark.sh
+++ b/benchmarks/disagg_benchmarks/disagg_overhead_benchmark.sh
@@ -10,7 +10,8 @@ set -ex
kill_gpu_processes() {
# kill all processes on GPU.
- pkill -f pt_main_thread
+ pgrep pt_main_thread | xargs -r kill -9
+ pgrep python3 | xargs -r kill -9
sleep 10
# remove vllm config file
@@ -54,7 +55,7 @@ benchmark() {
CUDA_VISIBLE_DEVICES=0 python3 \
-m vllm.entrypoints.openai.api_server \
- --model meta-llama/Meta-Llama-3.1-8B-Instruct \
+ --model $model \
--port 8100 \
--max-model-len 10000 \
--gpu-memory-utilization 0.6 \
@@ -64,7 +65,7 @@ benchmark() {
CUDA_VISIBLE_DEVICES=1 python3 \
-m vllm.entrypoints.openai.api_server \
- --model meta-llama/Meta-Llama-3.1-8B-Instruct \
+ --model $model \
--port 8200 \
--max-model-len 10000 \
--gpu-memory-utilization 0.6 \
@@ -87,7 +88,7 @@ benchmark() {
--port 8100 \
--save-result \
--result-dir $results_folder \
- --result-filename disagg_prefill_2xtp4.json \
+ --result-filename disagg_prefill_tp1.json \
--request-rate "inf"
@@ -105,7 +106,7 @@ benchmark() {
--port 8200 \
--save-result \
--result-dir $results_folder \
- --result-filename disagg_prefill_2xtp4.json \
+ --result-filename disagg_prefill_tp1_overhead.json \
--request-rate "$qps"
kill_gpu_processes
@@ -118,7 +119,7 @@ main() {
(which jq) || (apt-get -y install jq)
(which socat) || (apt-get -y install socat)
- pip install quart httpx
+ pip install quart httpx datasets
cd "$(dirname "$0")"
diff --git a/benchmarks/disagg_benchmarks/disagg_performance_benchmark.sh b/benchmarks/disagg_benchmarks/disagg_performance_benchmark.sh
index d8d9e976dce76..eb5d891d0d4a5 100644
--- a/benchmarks/disagg_benchmarks/disagg_performance_benchmark.sh
+++ b/benchmarks/disagg_benchmarks/disagg_performance_benchmark.sh
@@ -1,13 +1,12 @@
#!/bin/bash
-# Requirement: 8x H100 GPUs.
+# Requirement: 2x GPUs.
-# Model: neuralmagic/Meta-Llama-3-70B-Instruct-FP8-KV
-# Query: 2048 input tokens, 11 output tokens, QPS 4, 500 requests
-# Resource: 8x H100
+# Model: meta-llama/Meta-Llama-3.1-8B-Instruct
+# Query: 1024 input tokens, 6 output tokens, QPS 2/4/6/8, 100 requests
+# Resource: 2x GPU
# Approaches:
-# 1. Chunked prefill: 1 vllm instance with tp=8
# 2. Chunked prefill: 2 vllm instance with tp=4, equivalent to 1 tp=4 instance with QPS 4
# 3. Disaggregated prefill: 1 prefilling instance and 1 decoding instance
# Prefilling instance: max_output_token=1
@@ -114,7 +113,6 @@ benchmark() {
--request-rate "$qps"
sleep 2
-
}
@@ -123,8 +121,9 @@ main() {
(which wget && which curl) || (apt-get update && apt-get install -y wget curl)
(which jq) || (apt-get -y install jq)
(which socat) || (apt-get -y install socat)
+ (which lsof) || (apt-get -y install lsof)
- pip install quart httpx matplotlib aiohttp
+ pip install quart httpx matplotlib aiohttp datasets
cd "$(dirname "$0")"
diff --git a/benchmarks/fused_kernels/layernorm_rms_benchmarks.py b/benchmarks/fused_kernels/layernorm_rms_benchmarks.py
new file mode 100644
index 0000000000000..ef91f9f8eb529
--- /dev/null
+++ b/benchmarks/fused_kernels/layernorm_rms_benchmarks.py
@@ -0,0 +1,173 @@
+import pickle as pkl
+import time
+from dataclasses import dataclass
+from itertools import product
+from typing import Callable, Iterable, List, Optional
+
+import torch
+import torch.utils.benchmark as TBenchmark
+from torch.utils.benchmark import Measurement as TMeasurement
+from tqdm import tqdm
+
+import vllm._custom_ops as ops
+from vllm.model_executor.layers.layernorm import RMSNorm
+
+
+@dataclass
+class bench_params_t:
+ num_tokens: int
+ hidden_size: int
+ add_residual: bool
+ dtype: torch.dtype
+
+ def description(self):
+ return (f'N {self.num_tokens} '
+ f'x D {self.hidden_size} '
+ f'x R {self.add_residual} '
+ f'x DT {self.dtype}')
+
+
+def get_bench_params() -> List[bench_params_t]:
+ ## Test Fixtures
+ NUM_TOKENS = [2**x for x in range(11)]
+ HIDDEN_SIZES = list(range(1024, 8129, 1024))
+ ADD_RESIDUAL = [True, False]
+ DTYPES = [torch.bfloat16, torch.float]
+
+ combinations = product(NUM_TOKENS, HIDDEN_SIZES, ADD_RESIDUAL, DTYPES)
+ bench_params = list(map(lambda x: \
+ bench_params_t(x[0], x[1], x[2], x[3]), combinations))
+ return bench_params
+
+
+# Reference impls
+def unfused_int8_impl(rms_norm_layer: RMSNorm, x: torch.Tensor,
+ residual: Optional[torch.Tensor],
+ quant_dtype: torch.dtype):
+ # Norm
+ torch_out = None
+ if residual is None:
+ torch_out = rms_norm_layer.forward_cuda(x, residual)
+ else:
+ torch_out, _ = rms_norm_layer.forward_cuda(x, residual)
+
+ # Quant
+ torch_out, _, _ = ops.scaled_int8_quant(torch_out)
+
+
+def unfused_fp8_impl(rms_norm_layer: RMSNorm, x: torch.Tensor,
+ residual: Optional[torch.Tensor],
+ quant_dtype: torch.dtype):
+ # Norm
+ torch_out = None
+ if residual is None:
+ torch_out = rms_norm_layer.forward_cuda(x, residual)
+ else:
+ torch_out, _ = rms_norm_layer.forward_cuda(x, residual)
+
+ # Quant
+ torch_out, _ = ops.scaled_fp8_quant(torch_out)
+
+
+def fused_impl(
+ rms_norm_layer: RMSNorm, # this stores the weights
+ x: torch.Tensor,
+ residual: Optional[torch.Tensor],
+ quant_dtype: torch.dtype):
+ out, _ = ops.rms_norm_dynamic_per_token_quant(x,
+ rms_norm_layer.weight,
+ 1e-6,
+ quant_dtype,
+ residual=residual)
+
+
+# Bench functions
+def bench_fn(rms_norm_layer: RMSNorm, x: torch.Tensor, residual: torch.Tensor,
+ quant_dtype: torch.dtype, label: str, sub_label: str,
+ fn: Callable, description: str) -> TMeasurement:
+
+ min_run_time = 1
+
+ globals = {
+ "rms_norm_layer": rms_norm_layer,
+ "x": x,
+ "residual": residual,
+ "quant_dtype": quant_dtype,
+ "fn": fn,
+ }
+ return TBenchmark.Timer(
+ stmt="fn(rms_norm_layer, x, residual, quant_dtype)",
+ globals=globals,
+ label=label,
+ sub_label=sub_label,
+ description=description,
+ ).blocked_autorange(min_run_time=min_run_time)
+
+def bench(params: bench_params_t, label: str, sub_label: str) \
+ -> Iterable[TMeasurement]:
+
+ # Make inputs
+ layer = RMSNorm(params.hidden_size, 1e-6).to(dtype=params.dtype)
+ # Make weights
+ layer.weight.data.normal_(mean=1.0, std=0.1)
+ # Make inputs
+ scale = 1 / params.hidden_size
+ x = torch.randn(params.num_tokens,
+ params.hidden_size,
+ dtype=params.dtype,
+ device='cuda') * scale
+ residual = (torch.randn_like(x) * scale).to(device='cuda') \
+ if params.add_residual else None
+
+ timers = []
+
+ # unfused int8 impl.
+ timers.append(
+ bench_fn(layer, x, residual, torch.int8, label, sub_label,
+ unfused_int8_impl, "unfused_int8_impl"))
+
+ # unfused fp8 impl.
+ timers.append(
+ bench_fn(layer, x, residual, torch.float8_e4m3fn, label, sub_label,
+ unfused_fp8_impl, "unfused_fp8_impl"))
+
+ # fused int8 impl.
+ timers.append(
+ bench_fn(layer, x, residual, torch.int8, label, sub_label, fused_impl,
+ "fused_int8_impl"))
+
+ # fused fp8 impl.
+ timers.append(
+ bench_fn(layer, x, residual, torch.float8_e4m3fn, label, sub_label,
+ fused_impl, "fused_fp8_impl"))
+
+ print_timers(timers)
+
+ return timers
+
+
+# launch bench
+# runner
+def print_timers(timers: Iterable[TMeasurement]):
+ compare = TBenchmark.Compare(timers)
+ compare.print()
+
+
+def main():
+ torch.set_default_device('cuda')
+ bench_params = get_bench_params()
+
+ timers = []
+ for bp in tqdm(bench_params):
+ timers.extend(
+ bench(bp, "rms-norm-dynamic-per-token-quant", bp.description()))
+ print_timers(timers)
+
+ # pickle all the results
+ timestamp = int(time.time())
+ with open(f"rms_norm_dpt_quant-{timestamp}.pkl", "wb") as f:
+ pkl.dump(timers, f)
+
+
+if __name__ == '__main__':
+ main()
diff --git a/benchmarks/kernels/benchmark_mixtral_moe_rocm.py b/benchmarks/kernels/benchmark_mixtral_moe_rocm.py
index e7e2df5dd3ebd..d6b3e3bdf91b8 100644
--- a/benchmarks/kernels/benchmark_mixtral_moe_rocm.py
+++ b/benchmarks/kernels/benchmark_mixtral_moe_rocm.py
@@ -9,7 +9,6 @@
import triton.language as tl
from tqdm import tqdm
-import vllm._moe_C as moe_kernels
from vllm import _custom_ops as ops
from vllm.model_executor.layers.fused_moe import (get_config_file_name,
invoke_fused_moe_kernel,
@@ -225,7 +224,7 @@ def run_timing(
)
w1 = torch.rand(
- (num_total_experts, 2 * shard_intermediate_size, d_model+128),
+ (num_total_experts, 2 * shard_intermediate_size, d_model + 128),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
@@ -326,8 +325,7 @@ def run_timing(
compute_type=(tl.bfloat16 if hidden_states.dtype == torch.bfloat16
else tl.float16),
use_fp8_w8a8=False,
- use_int8_w8a16=False
- )
+ use_int8_w8a16=False)
ops.silu_and_mul(intermediate_cache2, intermediate_cache1.view(-1, N))
@@ -348,8 +346,7 @@ def run_timing(
compute_type=(tl.bfloat16 if hidden_states.dtype == torch.bfloat16
else tl.float16),
use_fp8_w8a8=False,
- use_int8_w8a16=False
- )
+ use_int8_w8a16=False)
end_event.record()
end_event.synchronize()
diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py
index 673be4256e05c..8471541e14412 100644
--- a/benchmarks/kernels/benchmark_moe.py
+++ b/benchmarks/kernels/benchmark_moe.py
@@ -43,9 +43,9 @@ def benchmark_config(
x = torch.randn(num_tokens, hidden_size, dtype=dtype)
padding_size = 0
if envs.VLLM_MOE_PADDING and not (use_fp8_w8a8 or use_int8_w8a16):
- padding_size = 128 # fp16 padding size
+ padding_size = 128 # fp16 padding size
if envs.VLLM_FP8_PADDING and use_fp8_w8a8:
- padding_size = 256 # fp8 padding size. Ignoring int8 for now
+ padding_size = 256 # fp8 padding size. Ignoring int8 for now
if use_int8_w8a16:
w1 = torch.randint(-127,
diff --git a/benchmarks/kernels/benchmark_rmsnorm.py b/benchmarks/kernels/benchmark_rmsnorm.py
new file mode 100644
index 0000000000000..baa5de0fff1bd
--- /dev/null
+++ b/benchmarks/kernels/benchmark_rmsnorm.py
@@ -0,0 +1,262 @@
+import itertools
+from typing import Optional, Tuple, Union
+
+import torch
+import triton
+from flashinfer.norm import fused_add_rmsnorm, rmsnorm
+from torch import nn
+
+from vllm import _custom_ops as vllm_ops
+
+
+class HuggingFaceRMSNorm(nn.Module):
+
+ def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ residual: Optional[torch.Tensor] = None,
+ ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ orig_dtype = x.dtype
+ x = x.to(torch.float32)
+ if residual is not None:
+ x = x + residual.to(torch.float32)
+ residual = x.to(orig_dtype)
+
+ variance = x.pow(2).mean(dim=-1, keepdim=True)
+ x = x * torch.rsqrt(variance + self.variance_epsilon)
+ x = x.to(orig_dtype) * self.weight
+ if residual is None:
+ return x
+ else:
+ return x, residual
+
+
+def rmsnorm_naive(
+ x: torch.Tensor,
+ weight: torch.Tensor,
+ residual: Optional[torch.Tensor] = None,
+ eps: float = 1e-6,
+):
+ naive_norm = HuggingFaceRMSNorm(x.shape[-1], eps=eps)
+ naive_norm.weight = nn.Parameter(weight)
+ naive_norm = naive_norm.to(x.device)
+
+ orig_shape = x.shape
+ x = x.view(-1, x.shape[-1])
+ if residual is not None:
+ residual = residual.view(-1, residual.shape[-1])
+
+ output = naive_norm(x, residual)
+
+ if isinstance(output, tuple):
+ output = (output[0].view(orig_shape), output[1].view(orig_shape))
+ else:
+ output = output.view(orig_shape)
+ return output
+
+
+def rmsnorm_flashinfer(
+ x: torch.Tensor,
+ weight: torch.Tensor,
+ residual: Optional[torch.Tensor] = None,
+ eps: float = 1e-6,
+):
+ orig_shape = x.shape
+ x = x.view(-1, x.shape[-1])
+ if residual is not None:
+ residual = residual.view(-1, residual.shape[-1])
+
+ if residual is not None:
+ fused_add_rmsnorm(x, residual, weight, eps)
+ output = (x, residual)
+ else:
+ output = rmsnorm(x, weight, eps)
+
+ if isinstance(output, tuple):
+ output = (output[0].view(orig_shape), output[1].view(orig_shape))
+ else:
+ output = output.view(orig_shape)
+ return output
+
+
+def rmsnorm_vllm(
+ x: torch.Tensor,
+ weight: torch.Tensor,
+ residual: Optional[torch.Tensor] = None,
+ eps: float = 1e-6,
+):
+ orig_shape = x.shape
+ x = x.view(-1, x.shape[-1])
+ if residual is not None:
+ residual = residual.view(-1, residual.shape[-1])
+
+ if residual is not None:
+ vllm_ops.fused_add_rms_norm(x, residual, weight, eps)
+ output = (x, residual)
+ else:
+ out = torch.empty_like(x)
+ vllm_ops.rms_norm(out, x, weight, eps)
+ output = out
+
+ if isinstance(output, tuple):
+ output = (output[0].view(orig_shape), output[1].view(orig_shape))
+ else:
+ output = output.view(orig_shape)
+ return output
+
+
+def calculate_diff(batch_size, seq_len, hidden_size, use_residual=True):
+ dtype = torch.bfloat16
+ x = torch.randn(batch_size,
+ seq_len,
+ hidden_size,
+ dtype=dtype,
+ device="cuda")
+ weight = torch.ones(hidden_size, dtype=dtype, device="cuda")
+ residual = torch.randn_like(x) if use_residual else None
+
+ output_naive = rmsnorm_naive(
+ x.clone(), weight,
+ residual.clone() if residual is not None else None)
+ output_flashinfer = rmsnorm_flashinfer(
+ x.clone(), weight,
+ residual.clone() if residual is not None else None)
+ output_vllm = rmsnorm_vllm(
+ x.clone(), weight,
+ residual.clone() if residual is not None else None)
+
+ if use_residual:
+ output_naive = output_naive[0]
+ output_flashinfer = output_flashinfer[0]
+ output_vllm = output_vllm[0]
+
+ print(f"Naive output={output_naive}")
+ print(f"FlashInfer output={output_flashinfer}")
+ print(f"VLLM output={output_vllm}")
+
+ if torch.allclose(output_naive, output_flashinfer, atol=1e-2,
+ rtol=1e-2) and torch.allclose(
+ output_naive, output_vllm, atol=1e-2, rtol=1e-2):
+ print("✅ All implementations match")
+ else:
+ print("❌ Implementations differ")
+
+
+batch_size_range = [2**i for i in range(0, 7, 2)]
+seq_length_range = [2**i for i in range(6, 11, 1)]
+head_num_range = [32, 48]
+configs = list(
+ itertools.product(head_num_range, batch_size_range, seq_length_range))
+
+
+def get_benchmark(use_residual):
+
+ @triton.testing.perf_report(
+ triton.testing.Benchmark(
+ x_names=["head_num", "batch_size", "seq_len"],
+ x_vals=[list(_) for _ in configs],
+ line_arg="provider",
+ line_vals=["huggingface", "flashinfer", "vllm"],
+ line_names=["HuggingFace", "FlashInfer", "vLLM"],
+ styles=[("blue", "-"), ("green", "-"), ("red", "-")],
+ ylabel="us",
+ plot_name=
+ f"rmsnorm-perf-{'with' if use_residual else 'without'}-residual",
+ args={},
+ ))
+ def benchmark(head_num, batch_size, seq_len, provider):
+ dtype = torch.bfloat16
+ hidden_size = head_num * 128 # assuming head_dim = 128
+
+ x = torch.randn(batch_size,
+ seq_len,
+ hidden_size,
+ dtype=dtype,
+ device="cuda")
+ weight = torch.ones(hidden_size, dtype=dtype, device="cuda")
+ residual = torch.randn_like(x) if use_residual else None
+
+ quantiles = [0.5, 0.2, 0.8]
+
+ if provider == "huggingface":
+ ms, min_ms, max_ms = triton.testing.do_bench(
+ lambda: rmsnorm_naive(
+ x.clone(),
+ weight,
+ residual.clone() if residual is not None else None,
+ ),
+ quantiles=quantiles,
+ )
+ elif provider == "flashinfer":
+ ms, min_ms, max_ms = triton.testing.do_bench(
+ lambda: rmsnorm_flashinfer(
+ x.clone(),
+ weight,
+ residual.clone() if residual is not None else None,
+ ),
+ quantiles=quantiles,
+ )
+ else:
+ ms, min_ms, max_ms = triton.testing.do_bench(
+ lambda: rmsnorm_vllm(
+ x.clone(),
+ weight,
+ residual.clone() if residual is not None else None,
+ ),
+ quantiles=quantiles,
+ )
+
+ return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+ return benchmark
+
+
+if __name__ == "__main__":
+ import argparse
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--batch-size",
+ type=int,
+ default=4,
+ help="Batch size",
+ )
+ parser.add_argument(
+ "--seq-len",
+ type=int,
+ default=128,
+ help="Sequence length",
+ )
+ parser.add_argument(
+ "--hidden-size",
+ type=int,
+ default=4096,
+ help="Hidden size (2nd dimension) of the sequence",
+ )
+ parser.add_argument("--use-residual",
+ action="store_true",
+ help="Whether to use residual connection")
+ parser.add_argument(
+ "--save-path",
+ type=str,
+ default="./configs/rmsnorm/",
+ help="Path to save rmsnorm benchmark results",
+ )
+
+ args = parser.parse_args()
+
+ # Run correctness test
+ calculate_diff(batch_size=args.batch_size,
+ seq_len=args.seq_len,
+ hidden_size=args.hidden_size,
+ use_residual=args.use_residual)
+
+ # Get the benchmark function with proper use_residual setting
+ benchmark = get_benchmark(args.use_residual)
+ # Run performance benchmark
+ benchmark.run(print_data=True, save_path=args.save_path)
diff --git a/benchmarks/structured_schemas/structured_schema_1.json b/benchmarks/structured_schemas/structured_schema_1.json
new file mode 100644
index 0000000000000..6003698469e8d
--- /dev/null
+++ b/benchmarks/structured_schemas/structured_schema_1.json
@@ -0,0 +1,113 @@
+{
+ "$schema":
+ "https://json-schema.org/draft/2020-12/schema",
+ "title":
+ "User Profile",
+ "type":
+ "object",
+ "properties": {
+ "userId": {
+ "type": "string",
+ "description": "Unique identifier for the user."
+ },
+ "personalInfo": {
+ "type": "object",
+ "properties": {
+ "firstName": {
+ "type": "string",
+ "description": "The user's first name."
+ },
+ "lastName": {
+ "type": "string",
+ "description": "The user's last name."
+ },
+ "age": {
+ "type": "integer",
+ "minimum": 0,
+ "description": "The user's age."
+ },
+ "phoneNumbers": {
+ "type":
+ "array",
+ "items": {
+ "type": "object",
+ "properties": {
+ "type": {
+ "type": "string",
+ "enum": ["home", "work", "mobile"],
+ "description": "Type of phone number."
+ },
+ "number": {
+ "type": "string",
+ "pattern": "^\\+?[1-9]\\d{1,14}$",
+ "description": "Phone number in E.164 format."
+ }
+ },
+ "required": ["type", "number"]
+ },
+ "description":
+ "List of phone numbers associated with the user."
+ }
+ },
+ "required": ["firstName", "lastName"]
+ },
+ "address": {
+ "type": "object",
+ "properties": {
+ "street": {
+ "type": "string",
+ "description": "Street address."
+ },
+ "city": {
+ "type": "string",
+ "description": "City name."
+ },
+ "state": {
+ "type": "string",
+ "description": "State or province."
+ },
+ "postalCode": {
+ "type": "string",
+ "pattern": "^\\d{5}(-\\d{4})?$",
+ "description": "Postal code."
+ },
+ "country": {
+ "type": "string",
+ "description": "Country name."
+ }
+ },
+ "required": ["street", "city", "state", "postalCode", "country"]
+ },
+ "preferences": {
+ "type": "object",
+ "properties": {
+ "newsletterSubscribed": {
+ "type":
+ "boolean",
+ "description":
+ "Indicates if the user is subscribed to the newsletter."
+ },
+ "favoriteCategories": {
+ "type": "array",
+ "items": {
+ "type": "string"
+ },
+ "description": "List of user's favorite categories."
+ }
+ },
+ "required": ["newsletterSubscribed"]
+ },
+ "accountStatus": {
+ "type": "string",
+ "enum": ["active", "inactive", "suspended"],
+ "description": "Current status of the user's account."
+ },
+ "registrationDate": {
+ "type": "string",
+ "format": "date-time",
+ "description": "ISO 8601 formatted date-time of user registration."
+ }
+ },
+ "required":
+ ["userId", "personalInfo", "address", "accountStatus", "registrationDate"]
+}
\ No newline at end of file
diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake
index 68f7ca1af05ad..714abca2a5ff7 100644
--- a/cmake/cpu_extension.cmake
+++ b/cmake/cpu_extension.cmake
@@ -4,6 +4,11 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS ON)
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+if (${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
+ set(MACOSX_FOUND TRUE)
+endif()
+
+
#
# Define environment variables for special configurations
#
@@ -13,6 +18,9 @@ endif()
include_directories("${CMAKE_SOURCE_DIR}/csrc")
+
+set (ENABLE_NUMA TRUE)
+
#
# Check the compile flags
#
@@ -22,18 +30,28 @@ if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64")
"-mf16c"
)
endif()
-list(APPEND CXX_COMPILE_FLAGS
- "-fopenmp"
- "-DVLLM_CPU_EXTENSION")
-execute_process(COMMAND cat /proc/cpuinfo
- RESULT_VARIABLE CPUINFO_RET
- OUTPUT_VARIABLE CPUINFO)
+if(MACOSX_FOUND)
+ list(APPEND CXX_COMPILE_FLAGS
+ "-Xpreprocessor"
+ "-fopenmp"
+ "-DVLLM_CPU_EXTENSION")
+else()
+ list(APPEND CXX_COMPILE_FLAGS
+ "-fopenmp"
+ "-DVLLM_CPU_EXTENSION")
+endif()
-if (NOT CPUINFO_RET EQUAL 0)
- message(FATAL_ERROR "Failed to check CPU features via /proc/cpuinfo")
+if (NOT MACOSX_FOUND)
+ execute_process(COMMAND cat /proc/cpuinfo
+ RESULT_VARIABLE CPUINFO_RET
+ OUTPUT_VARIABLE CPUINFO)
+ if (NOT CPUINFO_RET EQUAL 0)
+ message(FATAL_ERROR "Failed to check CPU features via /proc/cpuinfo")
+ endif()
endif()
+
function (find_isa CPUINFO TARGET OUT)
string(FIND ${CPUINFO} ${TARGET} ISA_FOUND)
if(NOT ISA_FOUND EQUAL -1)
@@ -54,12 +72,17 @@ endfunction()
is_avx512_disabled(AVX512_DISABLED)
-find_isa(${CPUINFO} "avx2" AVX2_FOUND)
-find_isa(${CPUINFO} "avx512f" AVX512_FOUND)
-find_isa(${CPUINFO} "POWER10" POWER10_FOUND)
-find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
-find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
-find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
+if (MACOSX_FOUND AND CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64")
+ set(APPLE_SILICON_FOUND TRUE)
+else()
+ find_isa(${CPUINFO} "avx2" AVX2_FOUND)
+ find_isa(${CPUINFO} "avx512f" AVX512_FOUND)
+ find_isa(${CPUINFO} "POWER10" POWER10_FOUND)
+ find_isa(${CPUINFO} "POWER9" POWER9_FOUND)
+ find_isa(${CPUINFO} "asimd" ASIMD_FOUND) # Check for ARM NEON support
+ find_isa(${CPUINFO} "bf16" ARM_BF16_FOUND) # Check for ARM BF16 support
+endif()
+
if (AVX512_FOUND AND NOT AVX512_DISABLED)
list(APPEND CXX_COMPILE_FLAGS
@@ -103,6 +126,9 @@ elseif (ASIMD_FOUND)
set(MARCH_FLAGS "-march=armv8.2-a+dotprod+fp16")
endif()
list(APPEND CXX_COMPILE_FLAGS ${MARCH_FLAGS})
+elseif(APPLE_SILICON_FOUND)
+ message(STATUS "Apple Silicon Detected")
+ set(ENABLE_NUMA OFF)
else()
message(FATAL_ERROR "vLLM CPU backend requires AVX512, AVX2, Power9+ ISA or ARMv8 support.")
endif()
@@ -139,7 +165,12 @@ endif()
message(STATUS "CPU extension compile flags: ${CXX_COMPILE_FLAGS}")
-list(APPEND LIBS numa)
+if(ENABLE_NUMA)
+ list(APPEND LIBS numa)
+else()
+ message(STATUS "NUMA is disabled")
+ add_compile_definitions(-DVLLM_NUMA_DISABLED)
+endif()
#
# _C extension
diff --git a/csrc/attention/attention_kernels.cuh b/csrc/attention/attention_kernels.cuh
index 563e1438f0b01..eb216dc8baf10 100644
--- a/csrc/attention/attention_kernels.cuh
+++ b/csrc/attention/attention_kernels.cuh
@@ -105,7 +105,7 @@ __device__ void paged_attention_kernel(
const int max_num_blocks_per_seq,
const float* __restrict__ alibi_slopes, // [num_heads]
const int q_stride, const int kv_block_stride, const int kv_head_stride,
- const float k_scale, const float v_scale, const int tp_rank,
+ const float* k_scale, const float* v_scale, const int tp_rank,
const int blocksparse_local_blocks, const int blocksparse_vert_stride,
const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
const int seq_idx = blockIdx.y;
@@ -285,7 +285,7 @@ __device__ void paged_attention_kernel(
Quant_vec k_vec_quant = *reinterpret_cast(
k_ptr + offset1 * BLOCK_SIZE * x + offset2);
k_vecs[j] = fp8::scaled_convert(
- k_vec_quant, k_scale);
+ k_vec_quant, *k_scale);
}
}
@@ -415,7 +415,7 @@ __device__ void paged_attention_kernel(
*reinterpret_cast(v_ptr + offset);
// Vector conversion from V_quant_vec to V_vec.
v_vec = fp8::scaled_convert(v_quant_vec,
- v_scale);
+ *v_scale);
}
if (block_idx == num_seq_blocks - 1) {
// NOTE(woosuk): When v_vec contains the tokens that are out of the
@@ -513,7 +513,7 @@ __global__ void paged_attention_v1_kernel(
const int max_num_blocks_per_seq,
const float* __restrict__ alibi_slopes, // [num_heads]
const int q_stride, const int kv_block_stride, const int kv_head_stride,
- const float k_scale, const float v_scale, const int tp_rank,
+ const float* k_scale, const float* v_scale, const int tp_rank,
const int blocksparse_local_blocks, const int blocksparse_vert_stride,
const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
paged_attention_kernel& alibi_slopes, float k_scale,
- float v_scale, const int tp_rank, const int blocksparse_local_blocks,
- const int blocksparse_vert_stride, const int blocksparse_block_size,
- const int blocksparse_head_sliding_step) {
+ const std::optional& alibi_slopes, torch::Tensor& k_scale,
+ torch::Tensor& v_scale, const int tp_rank,
+ const int blocksparse_local_blocks, const int blocksparse_vert_stride,
+ const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
int num_seqs = query.size(0);
int num_heads = query.size(1);
int head_size = query.size(2);
@@ -80,6 +80,8 @@ void paged_attention_v1_launcher(
CACHE_T* value_cache_ptr = reinterpret_cast(value_cache.data_ptr());
int* block_tables_ptr = block_tables.data_ptr();
int* seq_lens_ptr = seq_lens.data_ptr();
+ const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr());
+ const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr());
constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
int padded_max_seq_len =
@@ -192,9 +194,10 @@ void paged_attention_v1(
torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
torch::Tensor& seq_lens, // [num_seqs]
int64_t block_size, int64_t max_seq_len,
- const c10::optional& alibi_slopes,
- const std::string& kv_cache_dtype, double k_scale, double v_scale,
- const int64_t tp_rank, const int64_t blocksparse_local_blocks,
+ const std::optional& alibi_slopes,
+ const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+ torch::Tensor& v_scale, const int64_t tp_rank,
+ const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
const int64_t blocksparse_head_sliding_step, const int64_t num_threads) {
const bool is_block_sparse = (blocksparse_vert_stride > 1);
diff --git a/csrc/attention/paged_attention_v2.cu b/csrc/attention/paged_attention_v2.cu
index 0c9caeb687ecb..bc543e713fe58 100644
--- a/csrc/attention/paged_attention_v2.cu
+++ b/csrc/attention/paged_attention_v2.cu
@@ -45,7 +45,7 @@ typedef __hip_bfloat16 __nv_bfloat16;
exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, key_cache_ptr, \
value_cache_ptr, num_kv_heads, scale, block_tables_ptr, \
seq_lens_ptr, max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, \
- kv_block_stride, kv_head_stride, k_scale, v_scale, tp_rank, \
+ kv_block_stride, kv_head_stride, k_scale_ptr, v_scale_ptr, tp_rank, \
blocksparse_local_blocks, blocksparse_vert_stride, \
blocksparse_block_size, blocksparse_head_sliding_step); \
vllm::paged_attention_v2_reduce_kernel& alibi_slopes, float k_scale,
- float v_scale, const int tp_rank, const int blocksparse_local_blocks,
- const int blocksparse_vert_stride, const int blocksparse_block_size,
- const int blocksparse_head_sliding_step) {
+ const std::optional& alibi_slopes, torch::Tensor& k_scale,
+ torch::Tensor& v_scale, const int tp_rank,
+ const int blocksparse_local_blocks, const int blocksparse_vert_stride,
+ const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
int num_seqs = query.size(0);
int num_heads = query.size(1);
int head_size = query.size(2);
@@ -92,6 +92,8 @@ void paged_attention_v2_launcher(
CACHE_T* value_cache_ptr = reinterpret_cast(value_cache.data_ptr());
int* block_tables_ptr = block_tables.data_ptr();
int* seq_lens_ptr = seq_lens.data_ptr();
+ const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr());
+ const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr());
constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
@@ -211,9 +213,10 @@ void paged_attention_v2(
torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
torch::Tensor& seq_lens, // [num_seqs]
int64_t block_size, int64_t max_seq_len,
- const c10::optional& alibi_slopes,
- const std::string& kv_cache_dtype, double k_scale, double v_scale,
- const int64_t tp_rank, const int64_t blocksparse_local_blocks,
+ const std::optional& alibi_slopes,
+ const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+ torch::Tensor& v_scale, const int64_t tp_rank,
+ const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
const int64_t blocksparse_head_sliding_step, const int64_t num_threads) {
const bool is_block_sparse = (blocksparse_vert_stride > 1);
diff --git a/csrc/cache.h b/csrc/cache.h
index 11c4c5001daaa..eedad9fafa3c0 100644
--- a/csrc/cache.h
+++ b/csrc/cache.h
@@ -18,15 +18,15 @@ void copy_blocks(std::vector const& key_caches,
void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
torch::Tensor& key_cache, torch::Tensor& value_cache,
torch::Tensor& slot_mapping,
- const std::string& kv_cache_dtype, const double k_scale,
- const double v_scale);
+ const std::string& kv_cache_dtype,
+ torch::Tensor& k_scale, torch::Tensor& v_scale);
void reshape_and_cache_flash(torch::Tensor& key, torch::Tensor& value,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& slot_mapping,
const std::string& kv_cache_dtype,
- const double k_scale, const double v_scale);
+ torch::Tensor& k_scale, torch::Tensor& v_scale);
// Just for unittest
void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
diff --git a/csrc/cache_kernels.cu b/csrc/cache_kernels.cu
index 1be806bbfa43c..21a0aec0ececc 100644
--- a/csrc/cache_kernels.cu
+++ b/csrc/cache_kernels.cu
@@ -159,8 +159,8 @@ __global__ void reshape_and_cache_kernel(
// block_size]
const int64_t* __restrict__ slot_mapping, // [num_tokens]
const int key_stride, const int value_stride, const int num_heads,
- const int head_size, const int block_size, const int x, const float k_scale,
- const float v_scale) {
+ const int head_size, const int block_size, const int x,
+ const float* k_scale, const float* v_scale) {
const int64_t token_idx = blockIdx.x;
const int64_t slot_idx = slot_mapping[token_idx];
if (slot_idx < 0) {
@@ -196,9 +196,9 @@ __global__ void reshape_and_cache_kernel(
value_cache[tgt_value_idx] = tgt_value;
} else {
key_cache[tgt_key_idx] =
- fp8::scaled_convert(tgt_key, k_scale);
+ fp8::scaled_convert(tgt_key, *k_scale);
value_cache[tgt_value_idx] =
- fp8::scaled_convert(tgt_value, v_scale);
+ fp8::scaled_convert(tgt_value, *v_scale);
}
}
}
@@ -214,7 +214,7 @@ __global__ void reshape_and_cache_flash_kernel(
const int64_t* __restrict__ slot_mapping, // [num_tokens]
const int block_stride, const int key_stride, const int value_stride,
const int num_heads, const int head_size, const int block_size,
- const float k_scale, const float v_scale) {
+ const float* k_scale, const float* v_scale) {
const int64_t token_idx = blockIdx.x;
const int64_t slot_idx = slot_mapping[token_idx];
// NOTE: slot_idx can be -1 if the token is padded
@@ -239,9 +239,9 @@ __global__ void reshape_and_cache_flash_kernel(
value_cache[tgt_key_value_idx] = tgt_value;
} else {
key_cache[tgt_key_value_idx] =
- fp8::scaled_convert(tgt_key, k_scale);
+ fp8::scaled_convert(tgt_key, *k_scale);
value_cache[tgt_key_value_idx] =
- fp8::scaled_convert(tgt_value, v_scale);
+ fp8::scaled_convert(tgt_value, *v_scale);
}
}
}
@@ -258,7 +258,9 @@ __global__ void reshape_and_cache_flash_kernel(
reinterpret_cast(key_cache.data_ptr()), \
reinterpret_cast(value_cache.data_ptr()), \
slot_mapping.data_ptr(), key_stride, value_stride, \
- num_heads, head_size, block_size, x, k_scale, v_scale);
+ num_heads, head_size, block_size, x, \
+ reinterpret_cast(k_scale.data_ptr()), \
+ reinterpret_cast(v_scale.data_ptr()));
void reshape_and_cache(
torch::Tensor& key, // [num_tokens, num_heads, head_size]
@@ -268,8 +270,8 @@ void reshape_and_cache(
torch::Tensor&
value_cache, // [num_blocks, num_heads, head_size, block_size]
torch::Tensor& slot_mapping, // [num_tokens]
- const std::string& kv_cache_dtype, const double k_scale,
- const double v_scale) {
+ const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+ torch::Tensor& v_scale) {
int num_tokens = key.size(0);
int num_heads = key.size(1);
int head_size = key.size(2);
@@ -299,7 +301,9 @@ void reshape_and_cache(
reinterpret_cast(key_cache.data_ptr()), \
reinterpret_cast(value_cache.data_ptr()), \
slot_mapping.data_ptr(), block_stride, key_stride, \
- value_stride, num_heads, head_size, block_size, k_scale, v_scale);
+ value_stride, num_heads, head_size, block_size, \
+ reinterpret_cast(k_scale.data_ptr()), \
+ reinterpret_cast(v_scale.data_ptr()));
void reshape_and_cache_flash(
torch::Tensor& key, // [num_tokens, num_heads, head_size]
@@ -307,10 +311,20 @@ void reshape_and_cache_flash(
torch::Tensor& key_cache, // [num_blocks, block_size, num_heads, head_size]
torch::Tensor&
value_cache, // [num_blocks, block_size, num_heads, head_size]
- torch::Tensor& slot_mapping, // [num_tokens]
- const std::string& kv_cache_dtype, const double k_scale,
- const double v_scale) {
- int num_tokens = key.size(0);
+ torch::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens]
+ const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+ torch::Tensor& v_scale) {
+ // NOTE(woosuk): In vLLM V1, key.size(0) can be different from
+ // slot_mapping.size(0) because of padding for CUDA graphs.
+ // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
+ // both include padding.
+ // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
+ // since key includes padding for CUDA graphs, while slot_mapping does not.
+ // In this case, slot_mapping.size(0) represents the actual number of tokens
+ // before padding.
+ // For compatibility with both cases, we use slot_mapping.size(0) as the
+ // number of tokens.
+ int num_tokens = slot_mapping.size(0);
int num_heads = key.size(1);
int head_size = key.size(2);
int block_size = key_cache.size(1);
diff --git a/csrc/core/math.hpp b/csrc/core/math.hpp
new file mode 100644
index 0000000000000..ba9f40a230c8e
--- /dev/null
+++ b/csrc/core/math.hpp
@@ -0,0 +1,7 @@
+#include
+#include
+
+inline uint32_t next_pow_2(uint32_t const num) {
+ if (num <= 1) return num;
+ return 1 << (CHAR_BIT * sizeof(num) - __builtin_clz(num - 1));
+}
\ No newline at end of file
diff --git a/csrc/cpu/attention.cpp b/csrc/cpu/attention.cpp
index e21832ba7582f..ef5b14088c63b 100644
--- a/csrc/cpu/attention.cpp
+++ b/csrc/cpu/attention.cpp
@@ -386,7 +386,7 @@ void paged_attention_v1_impl_launcher(
torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
torch::Tensor& value_cache, int num_kv_heads, float scale,
torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
- const c10::optional& alibi_slopes) {
+ const std::optional& alibi_slopes) {
int num_seqs = query.size(0);
int num_heads = query.size(1);
int head_size = query.size(2);
@@ -459,7 +459,7 @@ void paged_attention_v1(
torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
- int64_t max_seq_len, const c10::optional& alibi_slopes,
+ int64_t max_seq_len, const std::optional& alibi_slopes,
const std::string& kv_cache_dtype, double k_scale, double v_scale,
const int64_t tp_rank, const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
@@ -702,7 +702,7 @@ void paged_attention_v2_impl_launcher(
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
torch::Tensor& value_cache, int num_kv_heads, float scale,
torch::Tensor& block_tables, torch::Tensor& seq_lens, int block_size,
- int max_seq_len, const c10::optional& alibi_slopes) {
+ int max_seq_len, const std::optional& alibi_slopes) {
int num_seqs = query.size(0);
int num_heads = query.size(1);
int head_size = query.size(2);
@@ -781,7 +781,7 @@ void paged_attention_v2(
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
- int64_t max_seq_len, const c10::optional& alibi_slopes,
+ int64_t max_seq_len, const std::optional& alibi_slopes,
const std::string& kv_cache_dtype, double k_scale, double v_scale,
const int64_t tp_rank, const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
diff --git a/csrc/cpu/cpu_types_arm.hpp b/csrc/cpu/cpu_types_arm.hpp
index 73e0f8cb2e0fb..ae062a5b86892 100644
--- a/csrc/cpu/cpu_types_arm.hpp
+++ b/csrc/cpu/cpu_types_arm.hpp
@@ -91,11 +91,68 @@ struct FP16Vec16 : public Vec {
vst1q_f16(reinterpret_cast<__fp16*>(ptr) + 8, reg.val[1]);
}
}
+
+ // Note: below is the unrolled version of the following code:
+ //
+ // for (int i = 0; i < remainder; ++i) {
+ // reinterpret_cast<__fp16*>(ptr)[full_blocks * 8 + i] =
+ // vgetq_lane_f16(temp, i);
+ // }
+ //
+ // For macOS build (Clang), the arm/neon intrinsics function
+ // `vgetq_lane_f16` needs the parameter `i` to be constant at compile
+ // time.
if (remainder > 0) {
float16x8_t temp = reg.val[full_blocks];
- for (int i = 0; i < remainder; ++i) {
- reinterpret_cast<__fp16*>(ptr)[full_blocks * 8 + i] = vgetq_lane_f16(temp, i);
+ __fp16* fp16_ptr = reinterpret_cast<__fp16*>(ptr);
+ switch (remainder)
+ {
+ case 1:
+ fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
+ break;
+ case 2:
+ fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
+ fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
+ break;
+ case 3:
+ fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
+ fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
+ fp16_ptr[full_blocks * 8 + 2] = vgetq_lane_f16(temp, 2);
+ break;
+ case 4:
+ fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
+ fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
+ fp16_ptr[full_blocks * 8 + 2] = vgetq_lane_f16(temp, 2);
+ fp16_ptr[full_blocks * 8 + 3] = vgetq_lane_f16(temp, 3);
+ break;
+ case 5:
+ fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
+ fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
+ fp16_ptr[full_blocks * 8 + 2] = vgetq_lane_f16(temp, 2);
+ fp16_ptr[full_blocks * 8 + 3] = vgetq_lane_f16(temp, 3);
+ fp16_ptr[full_blocks * 8 + 4] = vgetq_lane_f16(temp, 4);
+ break;
+ case 6:
+ fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
+ fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
+ fp16_ptr[full_blocks * 8 + 2] = vgetq_lane_f16(temp, 2);
+ fp16_ptr[full_blocks * 8 + 3] = vgetq_lane_f16(temp, 3);
+ fp16_ptr[full_blocks * 8 + 4] = vgetq_lane_f16(temp, 4);
+ fp16_ptr[full_blocks * 8 + 5] = vgetq_lane_f16(temp, 5);
+ break;
+ case 7:
+ fp16_ptr[full_blocks * 8 + 0] = vgetq_lane_f16(temp, 0);
+ fp16_ptr[full_blocks * 8 + 1] = vgetq_lane_f16(temp, 1);
+ fp16_ptr[full_blocks * 8 + 2] = vgetq_lane_f16(temp, 2);
+ fp16_ptr[full_blocks * 8 + 3] = vgetq_lane_f16(temp, 3);
+ fp16_ptr[full_blocks * 8 + 4] = vgetq_lane_f16(temp, 4);
+ fp16_ptr[full_blocks * 8 + 5] = vgetq_lane_f16(temp, 5);
+ fp16_ptr[full_blocks * 8 + 6] = vgetq_lane_f16(temp, 6);
+ break;
+
+ default:
+ break;
}
}
}
diff --git a/csrc/cpu/quant.cpp b/csrc/cpu/quant.cpp
index d9aed657a3113..33b1637832888 100644
--- a/csrc/cpu/quant.cpp
+++ b/csrc/cpu/quant.cpp
@@ -359,7 +359,7 @@ void int8_scaled_mm(torch::Tensor& c, // [M, OC], row-major
const torch::Tensor& b, // [IC, OC], column-major
const torch::Tensor& a_scales, // [1] or [M]
const torch::Tensor& b_scales, // [1] or [OC]
- const c10::optional& bias // [OC]
+ const std::optional& bias // [OC]
) {
CPU_KERNEL_GUARD_IN(cutlass_scaled_mm)
// Checks for conformality
@@ -442,8 +442,8 @@ void int8_scaled_mm_azp(torch::Tensor& c, // [M, OC], row-major
const torch::Tensor& a_scales, // [1] or [M]
const torch::Tensor& b_scales, // [1] or [OC]
const torch::Tensor& azp_adj, // [OC]
- const c10::optional& azp, // [1] or [M]
- const c10::optional& bias // [OC]
+ const std::optional& azp, // [1] or [M]
+ const std::optional& bias // [OC]
) {
CPU_KERNEL_GUARD_IN(cutlass_scaled_mm_azp)
// Checks for conformality
@@ -561,7 +561,7 @@ void int8_scaled_mm_azp(torch::Tensor& c, // [M, OC], row-major
void static_scaled_int8_quant(torch::Tensor& out, // [..., hidden_size]
const torch::Tensor& input, // [..., hidden_size]
const torch::Tensor& scale,
- c10::optional const& azp) {
+ std::optional const& azp) {
CPU_KERNEL_GUARD_IN(static_scaled_int8_quant)
TORCH_CHECK(input.is_contiguous());
TORCH_CHECK(out.is_contiguous());
@@ -590,7 +590,7 @@ void dynamic_scaled_int8_quant(
torch::Tensor& out, // [..., hidden_size]
const torch::Tensor& input, // [..., hidden_size]
torch::Tensor& scale, // [..., 1]
- c10::optional const& azp) {
+ std::optional const& azp) {
CPU_KERNEL_GUARD_IN(dynamic_scaled_int8_quant)
TORCH_CHECK(input.is_contiguous());
TORCH_CHECK(out.is_contiguous());
diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp
index 03beefbc6de7d..74e4d8189d403 100644
--- a/csrc/cpu/torch_bindings.cpp
+++ b/csrc/cpu/torch_bindings.cpp
@@ -9,14 +9,14 @@ std::string init_cpu_threads_env(const std::string& cpu_ids);
void int8_scaled_mm(torch::Tensor& c, const torch::Tensor& a,
const torch::Tensor& b, const torch::Tensor& a_scales,
const torch::Tensor& b_scales,
- const c10::optional& bias);
+ const std::optional& bias);
void int8_scaled_mm_azp(torch::Tensor& c, const torch::Tensor& a,
const torch::Tensor& b, const torch::Tensor& a_scales,
const torch::Tensor& b_scales,
const torch::Tensor& azp_adj,
- const c10::optional& azp,
- const c10::optional& bias);
+ const std::optional& azp,
+ const std::optional& bias);
TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// vLLM custom ops
diff --git a/csrc/cpu/utils.cpp b/csrc/cpu/utils.cpp
index 1138a55df2f05..42a1c1d924bac 100644
--- a/csrc/cpu/utils.cpp
+++ b/csrc/cpu/utils.cpp
@@ -1,10 +1,22 @@
-#include
-#include
-#include
-#include
+#ifndef VLLM_NUMA_DISABLED
+ #include
+ #include
+ #include
+ #include
+#endif
#include "cpu_types.hpp"
+#ifdef VLLM_NUMA_DISABLED
+std::string init_cpu_threads_env(const std::string& cpu_ids) {
+ return std::string(
+ "Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has "
+ "no effect to setup thread affinity.");
+}
+
+#endif
+
+#ifndef VLLM_NUMA_DISABLED
std::string init_cpu_threads_env(const std::string& cpu_ids) {
bitmask* omp_cpu_mask = numa_parse_cpustring(cpu_ids.c_str());
TORCH_CHECK(omp_cpu_mask->size > 0);
@@ -57,7 +69,7 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
omp_lock_t writelock;
omp_init_lock(&writelock);
-#pragma omp parallel for schedule(static, 1)
+ #pragma omp parallel for schedule(static, 1)
for (size_t i = 0; i < omp_cpu_ids.size(); ++i) {
cpu_set_t mask;
CPU_ZERO(&mask);
@@ -88,3 +100,4 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {
return ss.str();
}
+#endif
\ No newline at end of file
diff --git a/csrc/cutlass_extensions/common.cpp b/csrc/cutlass_extensions/common.cpp
new file mode 100644
index 0000000000000..3d2093ab94297
--- /dev/null
+++ b/csrc/cutlass_extensions/common.cpp
@@ -0,0 +1,11 @@
+#include "cutlass_extensions/common.hpp"
+
+int32_t get_sm_version_num() {
+ int32_t major_capability, minor_capability;
+ cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor,
+ 0);
+ cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor,
+ 0);
+ int32_t version_num = major_capability * 10 + minor_capability;
+ return version_num;
+}
\ No newline at end of file
diff --git a/csrc/cutlass_extensions/common.hpp b/csrc/cutlass_extensions/common.hpp
new file mode 100644
index 0000000000000..85e359aa57113
--- /dev/null
+++ b/csrc/cutlass_extensions/common.hpp
@@ -0,0 +1,35 @@
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include
+#include "cuda_runtime.h"
+#include
+
+/**
+ * Helper function for checking CUTLASS errors
+ */
+#define CUTLASS_CHECK(status) \
+ { \
+ cutlass::Status error = status; \
+ TORCH_CHECK(error == cutlass::Status::kSuccess, \
+ cutlassGetStatusString(error)); \
+ }
+
+/**
+ * Panic wrapper for unwinding CUDA runtime errors
+ */
+#define CUDA_CHECK(status) \
+ { \
+ cudaError_t error = status; \
+ TORCH_CHECK(error == cudaSuccess, cudaGetErrorString(error)); \
+ }
+
+inline int get_cuda_max_shared_memory_per_block_opt_in(int const device) {
+ int max_shared_mem_per_block_opt_in = 0;
+ cudaDeviceGetAttribute(&max_shared_mem_per_block_opt_in,
+ cudaDevAttrMaxSharedMemoryPerBlockOptin,
+ device);
+ return max_shared_mem_per_block_opt_in;
+}
+
+int32_t get_sm_version_num();
diff --git a/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
index c69e87999ae71..ef413e6dd75c5 100644
--- a/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
+++ b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
@@ -1,3 +1,5 @@
+#pragma once
+
#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c2x.hpp"
/*
@@ -66,7 +68,7 @@ struct ScaledEpilogueBase {
// This overload handles the case where there might not be a tensor, in which
// case a nullptr is passed and a constant (0) is used.
template
- static auto args_from_tensor(c10::optional const& tensor) {
+ static auto args_from_tensor(std::optional const& tensor) {
static_assert(std::is_same_v>);
using Arguments = typename Descriptor::Arguments;
auto* data_ptr = tensor ? static_cast(tensor->data_ptr()) : nullptr;
@@ -221,7 +223,7 @@ struct ScaledEpilogueBiasAzp
static ArgumentType prepare_args(torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
torch::Tensor const& azp_adj,
- c10::optional const& bias) {
+ std::optional const& bias) {
auto a_args = SUPER::template args_from_tensor(a_scales);
auto b_args = SUPER::template args_from_tensor(b_scales);
auto bias_args = SUPER::template args_from_tensor(bias);
@@ -299,7 +301,7 @@ struct ScaledEpilogueBiasAzpToken
torch::Tensor const& b_scales,
torch::Tensor const& azp_adj,
torch::Tensor const& azp,
- c10::optional const& bias) {
+ std::optional const& bias) {
auto a_args = SUPER::template args_from_tensor(a_scales);
auto b_args = SUPER::template args_from_tensor(b_scales);
auto bias_args = SUPER::template args_from_tensor(bias);
diff --git a/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
index 95764ecddc79f..c590c66a66652 100644
--- a/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
+++ b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
@@ -1,3 +1,5 @@
+#pragma once
+
#include "cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp"
/*
@@ -36,13 +38,13 @@ struct ScaledEpilogueBase {
// Don't want to support nullptr by default
template
using ColLoad = cutlass::epilogue::fusion::Sm90ColBroadcast<
- 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T,
+ 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T, T,
Stride, Int<0>, Int<0>>, 128 / sizeof_bits_v, EnableNullPtr>;
// Don't want to support nullptr by default
template
using RowLoad = cutlass::epilogue::fusion::Sm90RowBroadcast<
- 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T,
+ 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T, T,
Stride, Int<1>, Int<0>>, 128 / sizeof_bits_v, EnableNullPtr>;
// This utility function constructs the arguments for the load descriptors
@@ -65,7 +67,7 @@ struct ScaledEpilogueBase {
// This overload handles the case where there might not be a tensor, in which
// case a nullptr is passed and a constant (0) is used.
template
- static auto args_from_tensor(c10::optional const& tensor) {
+ static auto args_from_tensor(std::optional const& tensor) {
using Arguments = typename Descriptor::Arguments;
auto* data_ptr = tensor ? static_cast(tensor->data_ptr()) : nullptr;
static_assert(std::is_same_v> ||
@@ -221,7 +223,7 @@ struct ScaledEpilogueBiasAzp
static ArgumentType prepare_args(torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
torch::Tensor const& azp_adj,
- c10::optional const& bias) {
+ std::optional const& bias) {
auto a_args = SUPER::template args_from_tensor(a_scales);
auto b_args = SUPER::template args_from_tensor(b_scales);
auto bias_args = SUPER::template args_from_tensor(bias);
@@ -297,7 +299,7 @@ struct ScaledEpilogueBiasAzpToken
torch::Tensor const& b_scales,
torch::Tensor const& azp_adj,
torch::Tensor const& azp,
- c10::optional const& bias) {
+ std::optional const& bias) {
auto a_args = SUPER::template args_from_tensor(a_scales);
auto b_args = SUPER::template args_from_tensor(b_scales);
auto bias_args = SUPER::template args_from_tensor(bias);
diff --git a/csrc/cutlass_extensions/torch_utils.hpp b/csrc/cutlass_extensions/torch_utils.hpp
index 2c78572521eec..a1ff933cce63f 100644
--- a/csrc/cutlass_extensions/torch_utils.hpp
+++ b/csrc/cutlass_extensions/torch_utils.hpp
@@ -97,7 +97,7 @@ static inline auto make_cute_layout(torch::Tensor const& tensor,
template
static inline auto maybe_make_cute_layout(
- c10::optional const& tensor,
+ std::optional const& tensor,
std::string_view name = "tensor") {
using Layout = decltype(make_cute_layout(*tensor));
diff --git a/csrc/cutlass_extensions/vllm_cutlass_library_extension.py b/csrc/cutlass_extensions/vllm_cutlass_library_extension.py
index a5beea1a35e49..b401736c9824b 100644
--- a/csrc/cutlass_extensions/vllm_cutlass_library_extension.py
+++ b/csrc/cutlass_extensions/vllm_cutlass_library_extension.py
@@ -14,9 +14,9 @@ class VLLMDataType(enum.Enum):
class MixedInputKernelScheduleType(enum.Enum):
- TmaWarpSpecializedMixedInput = enum_auto()
- TmaWarpSpecializedPingpongMixedInput = enum_auto()
- TmaWarpSpecializedCooperativeMixedInput = enum_auto()
+ TmaWarpSpecialized = enum_auto()
+ TmaWarpSpecializedPingpong = enum_auto()
+ TmaWarpSpecializedCooperative = enum_auto()
VLLMDataTypeNames: Dict[Union[VLLMDataType, DataType], str] = {
@@ -68,11 +68,11 @@ class MixedInputKernelScheduleType(enum.Enum):
MixedInputKernelScheduleType, KernelScheduleType], str] = {
**KernelScheduleTag, # type: ignore
**{
- MixedInputKernelScheduleType.TmaWarpSpecializedMixedInput:
- "cutlass::gemm::KernelTmaWarpSpecializedMixedInput",
- MixedInputKernelScheduleType.TmaWarpSpecializedPingpongMixedInput:
- "cutlass::gemm::KernelTmaWarpSpecializedPingpongMixedInput",
- MixedInputKernelScheduleType.TmaWarpSpecializedCooperativeMixedInput:
- "cutlass::gemm::KernelTmaWarpSpecializedCooperativeMixedInput",
+ MixedInputKernelScheduleType.TmaWarpSpecialized:
+ "cutlass::gemm::KernelTmaWarpSpecialized",
+ MixedInputKernelScheduleType.TmaWarpSpecializedPingpong:
+ "cutlass::gemm::KernelTmaWarpSpecializedPingpong",
+ MixedInputKernelScheduleType.TmaWarpSpecializedCooperative:
+ "cutlass::gemm::KernelTmaWarpSpecializedCooperative",
}
}
diff --git a/csrc/dispatch_utils.h b/csrc/dispatch_utils.h
index a634e1c3d4886..3b477ba081997 100644
--- a/csrc/dispatch_utils.h
+++ b/csrc/dispatch_utils.h
@@ -14,6 +14,20 @@
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
+// TODO(luka/varun): use FP8_TYPE macro after refactoring
+#ifdef USE_CUDA_FP8_FORMAT
+ #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
+ AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
+ AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
+#else
+ #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
+ AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) \
+ AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
+#endif
+
+#define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
+ AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
+
#define VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
diff --git a/csrc/mamba/causal_conv1d/causal_conv1d.cu b/csrc/mamba/causal_conv1d/causal_conv1d.cu
index 498d069c05f0d..f0e5533bcae60 100644
--- a/csrc/mamba/causal_conv1d/causal_conv1d.cu
+++ b/csrc/mamba/causal_conv1d/causal_conv1d.cu
@@ -53,12 +53,12 @@ void set_conv_params_fwd(ConvParamsBase ¶ms,
const at::Tensor x,
const at::Tensor weight,
const at::Tensor out,
- const c10::optional& bias,
+ const std::optional& bias,
bool silu_activation,
int64_t pad_slot_id,
- const c10::optional& query_start_loc = std::nullopt,
- const c10::optional& cache_indices = std::nullopt,
- const c10::optional& has_initial_state = std::nullopt) {
+ const std::optional& query_start_loc = std::nullopt,
+ const std::optional& cache_indices = std::nullopt,
+ const std::optional& has_initial_state = std::nullopt) {
// Reset the parameters
memset(¶ms, 0, sizeof(params));
@@ -93,11 +93,11 @@ void set_conv_params_fwd(ConvParamsBase ¶ms,
void causal_conv1d_fwd(const at::Tensor &x, const at::Tensor &weight,
- const c10::optional &bias_,
- const c10::optional &conv_states,
- const c10::optional &query_start_loc,
- const c10::optional &cache_indices,
- const c10::optional &has_initial_state,
+ const std::optional &bias_,
+ const std::optional &conv_states,
+ const std::optional &query_start_loc,
+ const std::optional &cache_indices,
+ const std::optional &has_initial_state,
bool silu_activation,
// used to identify padding entries if cache_indices provided
// in case of padding, the kernel will return early
@@ -194,10 +194,10 @@ void causal_conv1d_fwd(const at::Tensor &x, const at::Tensor &weight,
void causal_conv1d_update(const at::Tensor &x,
const at::Tensor &conv_state,
const at::Tensor &weight,
- const c10::optional &bias_,
+ const std::optional &bias_,
bool silu_activation,
- const c10::optional &cache_seqlens_,
- const c10::optional &conv_state_indices_,
+ const std::optional &cache_seqlens_,
+ const std::optional &conv_state_indices_,
// used to identify padding entries if cache_indices provided
// in case of padding, the kernel will return early
int64_t pad_slot_id) {
@@ -424,7 +424,7 @@ void causal_conv1d_fwd_kernel(ConvParamsBase params) {
// and the one before it (chunk = n_chunks - 1 and chunk = n_chunks - 2),
// (which occurs when `final_state_position` is a non-positivie index)
// we load the correct data from smem_exchange from both chunks, the last chunk iteration and the one before it
- if (final_state_position < 0 && seqlen > kWidth){
+ if (conv_states != nullptr && final_state_position < 0 && seqlen > kWidth){
input_t vals_load[kNElts] = {0};
if ((chunk == n_chunks - 2) && (tidx == kNThreads - 1)){
// chunk = n_chunks - 2, a segment of the final state sits in the last index
diff --git a/csrc/mamba/mamba_ssm/selective_scan_fwd.cu b/csrc/mamba/mamba_ssm/selective_scan_fwd.cu
index 71624696338d0..bd0a34119c82b 100644
--- a/csrc/mamba/mamba_ssm/selective_scan_fwd.cu
+++ b/csrc/mamba/mamba_ssm/selective_scan_fwd.cu
@@ -402,14 +402,14 @@ void set_ssm_params_fwd(SSMParamsBase ¶ms,
const torch::Tensor out,
const torch::Tensor z,
const torch::Tensor out_z,
- const c10::optional& D,
- const c10::optional& delta_bias,
+ const std::optional& D,
+ const std::optional& delta_bias,
const torch::Tensor ssm_states,
bool has_z,
bool delta_softplus,
- const c10::optional& query_start_loc,
- const c10::optional& cache_indices,
- const c10::optional& has_initial_state,
+ const std::optional& query_start_loc,
+ const std::optional& cache_indices,
+ const std::optional& has_initial_state,
bool varlen,
int64_t pad_slot_id) {
@@ -504,13 +504,13 @@ void set_ssm_params_fwd(SSMParamsBase ¶ms,
void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
const torch::Tensor &A, const torch::Tensor &B, const torch::Tensor &C,
- const c10::optional &D_,
- const c10::optional &z_,
- const c10::optional &delta_bias_,
+ const std::optional &D_,
+ const std::optional &z_,
+ const std::optional &delta_bias_,
bool delta_softplus,
- const c10::optional &query_start_loc,
- const c10::optional &cache_indices,
- const c10::optional &has_initial_state,
+ const std::optional &query_start_loc,
+ const std::optional &cache_indices,
+ const std::optional &has_initial_state,
const torch::Tensor &ssm_states,
// used to identify padding entries if cache_indices provided
// in case of padding, the kernel will return early
diff --git a/csrc/moe/moe_align_sum_kernels.cu b/csrc/moe/moe_align_sum_kernels.cu
index dd90c38d9a721..16fccae403338 100644
--- a/csrc/moe/moe_align_sum_kernels.cu
+++ b/csrc/moe/moe_align_sum_kernels.cu
@@ -112,6 +112,91 @@ __global__ void moe_align_block_size_kernel(scalar_t* __restrict__ topk_ids,
}
}
+// TODO(simon): this is temporarily adapted from
+// https://github.com/sgl-project/sglang/commit/31548116a8dc8c6df7e146e0587335a59fc5b9d7
+// we did this to unblock Deepseek V3 but there should be a better
+// implementation to manage shared memory.
+template
+__global__ void moe_align_block_size_global_mem_kernel(
+ scalar_t* __restrict__ topk_ids, int32_t* sorted_token_ids,
+ int32_t* expert_ids, int32_t* total_tokens_post_pad, int32_t num_experts,
+ int32_t block_size, size_t numel, int32_t* tokens_cnts, int32_t* cumsum) {
+ const size_t tokens_per_thread = CEILDIV(numel, blockDim.x);
+ const size_t start_idx = threadIdx.x * tokens_per_thread;
+
+ for (int i = 0; i < num_experts; ++i) {
+ tokens_cnts[index(num_experts, threadIdx.x + 1, i)] = 0;
+ }
+
+ /**
+ * In the first step we compute token_cnts[thread_index + 1][expert_index],
+ * which counts how many tokens in the token shard of thread_index are
+ * assigned to expert expert_index.
+ */
+ for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+ ++tokens_cnts[index(num_experts, threadIdx.x + 1, topk_ids[i])];
+ }
+
+ __syncthreads();
+
+ // For each expert we accumulate the token counts from the different threads.
+ for (int eid = threadIdx.x; eid < num_experts; eid += blockDim.x) {
+ tokens_cnts[index(num_experts, 0, eid)] = 0;
+ for (int i = 1; i <= blockDim.x; ++i) {
+ tokens_cnts[index(num_experts, i, eid)] +=
+ tokens_cnts[index(num_experts, i - 1, eid)];
+ }
+ }
+
+ __syncthreads();
+
+ // We accumulate the token counts of all experts in thread 0.
+ if (threadIdx.x == 0) {
+ cumsum[0] = 0;
+ for (int i = 1; i <= num_experts; ++i) {
+ cumsum[i] = cumsum[i - 1] +
+ CEILDIV(tokens_cnts[index(num_experts, blockDim.x, i - 1)],
+ block_size) *
+ block_size;
+ }
+ *total_tokens_post_pad = cumsum[num_experts];
+ }
+
+ __syncthreads();
+
+ /**
+ * For each expert, each thread processes the tokens of the corresponding
+ * blocks and stores the corresponding expert_id for each block.
+ */
+ for (int eid = threadIdx.x; eid < num_experts; eid += blockDim.x) {
+ for (int i = cumsum[eid]; i < cumsum[eid + 1]; i += block_size) {
+ expert_ids[i / block_size] = eid;
+ }
+ }
+
+ /**
+ * Each thread processes a token shard, calculating the index of each token
+ * after sorting by expert number. Given the example topk_ids =
+ * [0,1,2,1,2,3,0,3,4] and block_size = 4, then the output would be [0, 6, *,
+ * *, 1, 3, *, *, 2, 4, *, *, 5, 7, *, *, 8, *, *, *], where * represents a
+ * padding value(preset in python).
+ */
+ for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+ int32_t expert_id = topk_ids[i];
+ /** The cumsum[expert_id] stores the starting index of the tokens that the
+ * expert with expert_id needs to process, and
+ * tokens_cnts[threadIdx.x][expert_id] stores the indices of the tokens
+ * processed by the expert with expert_id within the current thread's token
+ * shard.
+ */
+ int32_t rank_post_pad =
+ tokens_cnts[index(num_experts, threadIdx.x, expert_id)] +
+ cumsum[expert_id];
+ sorted_token_ids[rank_post_pad] = i;
+ ++tokens_cnts[index(num_experts, threadIdx.x, expert_id)];
+ }
+}
+
template
__global__ void moe_sum_kernel(
scalar_t* __restrict__ out, // [..., d]
diff --git a/csrc/ops.h b/csrc/ops.h
index 665c859ac2950..e9cc8d2e215e2 100644
--- a/csrc/ops.h
+++ b/csrc/ops.h
@@ -33,9 +33,10 @@ void paged_attention_v1(
torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
- int64_t max_seq_len, const c10::optional& alibi_slopes,
- const std::string& kv_cache_dtype, double k_scale, double v_scale,
- const int64_t tp_rank, const int64_t blocksparse_local_blocks,
+ int64_t max_seq_len, const std::optional& alibi_slopes,
+ const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+ torch::Tensor& v_scale, const int64_t tp_rank,
+ const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
const int64_t blocksparse_head_sliding_step, const int64_t num_threads);
@@ -44,9 +45,10 @@ void paged_attention_v2(
torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
torch::Tensor& value_cache, int64_t num_kv_heads, double scale,
torch::Tensor& block_tables, torch::Tensor& seq_lens, int64_t block_size,
- int64_t max_seq_len, const c10::optional& alibi_slopes,
- const std::string& kv_cache_dtype, double k_scale, double v_scale,
- const int64_t tp_rank, const int64_t blocksparse_local_blocks,
+ int64_t max_seq_len, const std::optional& alibi_slopes,
+ const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+ torch::Tensor& v_scale, const int64_t tp_rank,
+ const int64_t blocksparse_local_blocks,
const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
const int64_t blocksparse_head_sliding_step, const int64_t num_threads);
@@ -66,6 +68,14 @@ void fused_add_rms_norm_static_fp8_quant(torch::Tensor& out,
torch::Tensor& weight,
torch::Tensor& scale, double epsilon);
+void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
+ torch::Tensor const& input,
+ torch::Tensor const& weight,
+ torch::Tensor& scales,
+ double const epsilon,
+ std::optional scale_ub,
+ std::optional residual);
+
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
torch::Tensor& key, int64_t head_size,
torch::Tensor& cos_sin_cache, bool is_neox);
@@ -148,24 +158,35 @@ bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
void cutlass_scaled_mm(torch::Tensor& out, torch::Tensor const& a,
torch::Tensor const& b, torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
- c10::optional const& bias);
+ std::optional const& bias);
void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a,
torch::Tensor const& b,
torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
torch::Tensor const& azp_adj,
- c10::optional const& azp,
- c10::optional const& bias);
+ std::optional const& azp,
+ std::optional const& bias);
+
+bool cutlass_sparse_scaled_mm_supported(int64_t cuda_device_capability);
+
+void cutlass_scaled_sparse_mm(torch::Tensor& out, torch::Tensor const& a,
+ torch::Tensor const& b, torch::Tensor const& e,
+ torch::Tensor const& a_scales,
+ torch::Tensor const& b_scales,
+ std::optional const& bias);
+
+bool cutlass_sparse_compress_entry(torch::Tensor& a_compressed,
+ torch::Tensor& e, torch::Tensor const& a);
#endif
void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
torch::Tensor const& scale,
- c10::optional const& azp);
+ std::optional const& azp);
void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
torch::Tensor& scales,
- c10::optional const& azp);
+ std::optional const& azp);
torch::Tensor gptq_gemm(torch::Tensor a, torch::Tensor b_q_weight,
torch::Tensor b_gptq_qzeros,
@@ -182,34 +203,34 @@ void dynamic_scaled_fp8_quant(torch::Tensor& out, torch::Tensor const& input,
void dynamic_per_token_scaled_fp8_quant(
torch::Tensor& out, torch::Tensor const& input, torch::Tensor& scale,
- c10::optional const& scale_ub);
+ std::optional const& scale_ub);
void selective_scan_fwd(const torch::Tensor& u, const torch::Tensor& delta,
const torch::Tensor& A, const torch::Tensor& B,
const torch::Tensor& C,
- const c10::optional& D_,
- const c10::optional& z_,
- const c10::optional& delta_bias_,
+ const std::optional& D_,
+ const std::optional& z_,
+ const std::optional& delta_bias_,
bool delta_softplus,
- const c10::optional& query_start_loc,
- const c10::optional& cache_indices,
- const c10::optional& has_initial_state,
+ const std::optional& query_start_loc,
+ const std::optional& cache_indices,
+ const std::optional& has_initial_state,
const torch::Tensor& ssm_states, int64_t pad_slot_id);
void causal_conv1d_update(const at::Tensor& x, const at::Tensor& conv_state,
const at::Tensor& weight,
- const c10::optional& bias_,
+ const std::optional& bias_,
bool silu_activation,
- const c10::optional& cache_seqlens_,
- const c10::optional& conv_state_indices_,
+ const std::optional& cache_seqlens_,
+ const std::optional& conv_state_indices_,
int64_t pad_slot_id);
void causal_conv1d_fwd(const at::Tensor& x, const at::Tensor& weight,
- const c10::optional& bias_,
- const c10::optional& conv_states,
- const c10::optional& query_start_loc,
- const c10::optional& cache_indices,
- const c10::optional& has_initial_state,
+ const std::optional& bias_,
+ const std::optional& conv_states,
+ const std::optional& query_start_loc,
+ const std::optional& cache_indices,
+ const std::optional& has_initial_state,
bool silu_activation, int64_t pad_slot_id);
using fptr_t = int64_t;
diff --git a/csrc/quantization/compressed_tensors/int8_quant_kernels.cu b/csrc/quantization/compressed_tensors/int8_quant_kernels.cu
index e9987535bd3ea..e79785827189d 100644
--- a/csrc/quantization/compressed_tensors/int8_quant_kernels.cu
+++ b/csrc/quantization/compressed_tensors/int8_quant_kernels.cu
@@ -226,7 +226,7 @@ __global__ void dynamic_scaled_int8_azp_quant_kernel(
void static_scaled_int8_quant(torch::Tensor& out, // [..., hidden_size]
torch::Tensor const& input, // [..., hidden_size]
torch::Tensor const& scale,
- c10::optional const& azp) {
+ std::optional const& azp) {
TORCH_CHECK(input.is_contiguous());
TORCH_CHECK(out.is_contiguous());
TORCH_CHECK(scale.numel() == 1);
@@ -257,7 +257,7 @@ void static_scaled_int8_quant(torch::Tensor& out, // [..., hidden_size]
void dynamic_scaled_int8_quant(
torch::Tensor& out, // [..., hidden_size]
torch::Tensor const& input, // [..., hidden_size]
- torch::Tensor& scales, c10::optional const& azp) {
+ torch::Tensor& scales, std::optional const& azp) {
TORCH_CHECK(input.is_contiguous());
TORCH_CHECK(out.is_contiguous());
TORCH_CHECK(scales.is_contiguous());
diff --git a/csrc/quantization/cutlass_w8a8/common.hpp b/csrc/quantization/cutlass_w8a8/common.hpp
deleted file mode 100644
index bf04bb400790f..0000000000000
--- a/csrc/quantization/cutlass_w8a8/common.hpp
+++ /dev/null
@@ -1,27 +0,0 @@
-#pragma once
-
-#include "cutlass/cutlass.h"
-#include
-
-/**
- * Helper function for checking CUTLASS errors
- */
-#define CUTLASS_CHECK(status) \
- { \
- TORCH_CHECK(status == cutlass::Status::kSuccess, \
- cutlassGetStatusString(status)) \
- }
-
-inline uint32_t next_pow_2(uint32_t const num) {
- if (num <= 1) return num;
- return 1 << (CHAR_BIT * sizeof(num) - __builtin_clz(num - 1));
-}
-
-inline int get_cuda_max_shared_memory_per_block_opt_in(int const device) {
- int max_shared_mem_per_block_opt_in = 0;
- cudaDeviceGetAttribute(&max_shared_mem_per_block_opt_in,
- cudaDevAttrMaxSharedMemoryPerBlockOptin,
- device);
- return max_shared_mem_per_block_opt_in;
-}
-
diff --git a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu
index dbb72e8bbd3f5..865fef5aeea11 100644
--- a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu
+++ b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu
@@ -39,7 +39,7 @@ void cutlass_scaled_mm_sm75(torch::Tensor& out, torch::Tensor const& a,
torch::Tensor const& b,
torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
- c10::optional const& bias) {
+ std::optional const& bias) {
TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
if (bias) {
@@ -58,8 +58,8 @@ void cutlass_scaled_mm_azp_sm75(torch::Tensor& out, torch::Tensor const& a,
torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
torch::Tensor const& azp_adj,
- c10::optional const& azp,
- c10::optional const& bias) {
+ std::optional const& azp,
+ std::optional const& bias) {
TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
@@ -94,7 +94,7 @@ void cutlass_scaled_mm_sm80(torch::Tensor& out, torch::Tensor const& a,
torch::Tensor const& b,
torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
- c10::optional const& bias) {
+ std::optional const& bias) {
TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
if (bias) {
@@ -113,8 +113,8 @@ void cutlass_scaled_mm_azp_sm80(torch::Tensor& out, torch::Tensor const& a,
torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
torch::Tensor const& azp_adj,
- c10::optional const& azp,
- c10::optional const& bias) {
+ std::optional const& azp,
+ std::optional const& bias) {
TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
@@ -165,7 +165,7 @@ void cutlass_scaled_mm_sm89(torch::Tensor& out, torch::Tensor const& a,
torch::Tensor const& b,
torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
- c10::optional const& bias) {
+ std::optional const& bias) {
TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
if (bias) {
@@ -184,8 +184,8 @@ void cutlass_scaled_mm_azp_sm89(torch::Tensor& out, torch::Tensor const& a,
torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
torch::Tensor const& azp_adj,
- c10::optional const& azp,
- c10::optional const& bias) {
+ std::optional const& azp,
+ std::optional const& bias) {
TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
diff --git a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh
index d03242f44ab1d..f2fae4b66d651 100644
--- a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh
+++ b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cuh
@@ -21,15 +21,16 @@
#include "cutlass/epilogue/threadblock/fusion/visitors.hpp"
#include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h"
-#include "common.hpp"
+#include "core/math.hpp"
+#include "cutlass_extensions/common.hpp"
// clang-format on
using namespace cute;
/*
- Epilogue functions can be defined to post-process the output before it is
- written to GPU memory.
- Epilogues must contain a public type named EVTCompute of type Sm80EVT,
+ Epilogues defined in,
+ csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
+ must contain a public type named EVTCompute of type Sm80EVT,
as well as a static prepare_args function that constructs an
EVTCompute::Arguments struct.
*/
diff --git a/csrc/quantization/cutlass_w8a8/scaled_mm_c3x.cu b/csrc/quantization/cutlass_w8a8/scaled_mm_c3x.cu
index 33581a63d4c3d..e18d7d79e5b77 100644
--- a/csrc/quantization/cutlass_w8a8/scaled_mm_c3x.cu
+++ b/csrc/quantization/cutlass_w8a8/scaled_mm_c3x.cu
@@ -1,384 +1,18 @@
-// clang-format will break include orders
-// clang-format off
#include
#if defined CUDA_VERSION && CUDA_VERSION >= 12000
-#include
+ #include "scaled_mm_c3x_sm90_fp8_dispatch.cuh"
+ #include "scaled_mm_c3x_sm90_int8_dispatch.cuh"
-#include
-
-#include
-#include
-#include
-
-#include "cutlass/cutlass.h"
-
-#include "cute/tensor.hpp"
-#include "cute/atom/mma_atom.hpp"
-#include "cutlass/numeric_types.h"
-
-#include "cutlass/gemm/device/gemm_universal_adapter.h"
-#include "cutlass/gemm/kernel/gemm_universal.hpp"
-#include "cutlass/epilogue/collective/collective_builder.hpp"
-#include "cutlass/gemm/collective/collective_builder.hpp"
-
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
-#include "common.hpp"
-// clang-format on
-
-using namespace cute;
+ #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
using namespace vllm;
/*
This file defines quantized GEMM operations using the CUTLASS 3.x API, for
NVIDIA GPUs with sm90a (Hopper) or later.
-
- Epilogue functions can be defined to post-process the output before it is
- written to GPU memory.
- Epilogues must contain a public type named EVTCompute of type Sm90EVT,
- as well as a static prepare_args function that constructs an
- EVTCompute::Arguments struct.
*/
-namespace {
-
-// A wrapper for the GEMM kernel that is used to guard against compilation on
-// architectures that will never use the kernel. The purpose of this is to
-// reduce the size of the compiled binary.
-// __CUDA_ARCH__ is not defined in host code, so this lets us smuggle the ifdef
-// into code that will be executed on the device where it is defined.
-template
-struct enable_sm90_or_later : Kernel {
- template
- CUTLASS_DEVICE void operator()(Args&&... args) {
- #if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 900
- Kernel::operator()(std::forward(args)...);
- #endif
- }
-};
-template typename Epilogue_,
- typename TileShape, typename ClusterShape, typename KernelSchedule,
- typename EpilogueSchedule>
-struct cutlass_3x_gemm {
- using ElementAB = ElementAB_;
- using ElementD = ElementD_;
- using ElementAcc =
- typename std::conditional, int32_t,
- float>::type;
-
- using EpilogueDescriptor =
- cutlass::epilogue::collective::detail::EpilogueDescriptor<
- TileShape, cutlass::epilogue::collective::EpilogueTileAuto, ElementD,
- ElementD, EpilogueSchedule>;
-
- using Epilogue = Epilogue_;
-
- using StrideD = Stride