forked from Karylab-cklius/vllm
Compare commits
114
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
30a3a7a7ff | ||
|
|
abe41f28de | ||
|
|
c3284c31f5 | ||
|
|
c74e751824 | ||
|
|
b93cbd7416 | ||
|
|
bdc6f3bfa1 | ||
|
|
392d1b4d2e | ||
|
|
21b396abe1 | ||
|
|
bdaf27519f | ||
|
|
beb4327c46 | ||
|
|
c46ced1ee3 | ||
|
|
65dcde1695 | ||
|
|
65a7b46284 | ||
|
|
93e2ab7111 | ||
|
|
8b745527cd | ||
|
|
920469974a | ||
|
|
8b91cd5b20 | ||
|
|
dd94484577 | ||
|
|
7ff656cc8b | ||
|
|
0a2965b1b3 | ||
|
|
0ed05b6f82 | ||
|
|
ed051fab54 | ||
|
|
48fcfc926c | ||
|
|
3354dba381 | ||
|
|
cbb5f045be | ||
|
|
b3e85be663 | ||
|
|
d3e69fd671 | ||
|
|
c85d72076a | ||
|
|
c5b66233b2 | ||
|
|
066f02ae94 | ||
|
|
5d23ca47ab | ||
|
|
e55cc59e52 | ||
|
|
ba50b9763f | ||
|
|
b4cfbc24d3 | ||
|
|
1e823dc01d | ||
|
|
8e61b646e2 | ||
|
|
e040899a00 | ||
|
|
dd5c299fbe | ||
|
|
6db31c8e76 | ||
|
|
cbe9c40f99 | ||
|
|
2f71b2bd9f | ||
|
|
32ab064621 | ||
|
|
c64c356990 | ||
|
|
34e6dfced8 | ||
|
|
39a1d32b59 | ||
|
|
700e882eab | ||
|
|
a4f019fa25 | ||
|
|
9dd2465896 | ||
|
|
a46c9329e5 | ||
|
|
445321fab4 | ||
|
|
69f3150981 | ||
|
|
86db6c3070 | ||
|
|
5769a7382c | ||
|
|
8484ca5d45 | ||
|
|
482e5524fe | ||
|
|
567a78432d | ||
|
|
d891b9bd51 | ||
|
|
04adc8843b | ||
|
|
ae098abe3f | ||
|
|
b1384f5ec6 | ||
|
|
b136cc2c2c | ||
|
|
9fde043f54 | ||
|
|
24dd2aec81 | ||
|
|
3ee9eea928 | ||
|
|
5bce653e09 | ||
|
|
5ad11172b7 | ||
|
|
f70caef48b | ||
|
|
8d8ec38361 | ||
|
|
b1c6dba558 | ||
|
|
598d51153a | ||
|
|
095adf1fdc | ||
|
|
51ee564e56 | ||
|
|
373eb314af | ||
|
|
641cb59592 | ||
|
|
07f9baf756 | ||
|
|
7a90eb98ab | ||
|
|
8f4c69b222 | ||
|
|
8b79971bb9 | ||
|
|
f676808ba0 | ||
|
|
98e4726a14 | ||
|
|
740f379fae | ||
|
|
40cc2e8327 | ||
|
|
ba22152096 | ||
|
|
90ce3a09be | ||
|
|
26c754d847 | ||
|
|
3d7f357ebf | ||
|
|
736f1a5907 | ||
|
|
344609ab17 | ||
|
|
d039c17114 | ||
|
|
cdab28319f | ||
|
|
2fa10566e3 | ||
|
|
fb265fc8fb | ||
|
|
8f0e75e16b | ||
|
|
98ba9b9583 | ||
|
|
990c2a0187 | ||
|
|
e433634c78 | ||
|
|
16f8110935 | ||
|
|
d9c1767cd4 | ||
|
|
e9cc1fd093 | ||
|
|
f1073c050c | ||
|
|
394edc8108 | ||
|
|
69715823df | ||
|
|
6569df6a3e | ||
|
|
f2aaf59151 | ||
|
|
95a248faed | ||
|
|
d2ec433e37 | ||
|
|
78a04c208d | ||
|
|
b71218107f | ||
|
|
cc1d020d01 | ||
|
|
8974ed89cd | ||
|
|
fb2faceacd | ||
|
|
b6cc46ec3b | ||
|
|
fa4321de3d | ||
|
|
9226613043 |
@@ -8,6 +8,7 @@ run_all_patterns:
|
||||
- "docker/docker-bake-rocm.hcl"
|
||||
- ".buildkite/hardware_tests/amd.yaml"
|
||||
- ".buildkite/scripts/ci-bake-rocm.sh"
|
||||
- ".buildkite/scripts/rocm/"
|
||||
- ".buildkite/scripts/hardware_ci/run-amd-test.py"
|
||||
- ".buildkite/scripts/hardware_ci/run-amd-test.sh"
|
||||
- "CMakeLists.txt"
|
||||
|
||||
@@ -1,5 +1,28 @@
|
||||
group: Hardware - AMD Build
|
||||
|
||||
# ROCm image flow:
|
||||
# 1. Refresh the long-lived ROCm base image only when Dockerfile.rocm_base changes.
|
||||
# 2. Build ci_base from either the stable base or the freshly refreshed base.
|
||||
# 3. Build the per-commit ROCm CI image and smoke-test it before GPU jobs run.
|
||||
steps:
|
||||
- label: "AMD: :docker: refresh ROCm base"
|
||||
key: refresh-rocm-base-amd
|
||||
depends_on: []
|
||||
device: amd_cpu
|
||||
no_plugin: true
|
||||
commands:
|
||||
- bash .buildkite/scripts/rocm/refresh-base-image.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
BUILDKIT_PROGRESS: "tty"
|
||||
TERM: "xterm-256color"
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: -1 # Agent was lost
|
||||
limit: 1
|
||||
- exit_status: -10 # Agent was lost
|
||||
limit: 1
|
||||
|
||||
# Ensure ci_base is up-to-date before building the test image.
|
||||
# Compares a content hash of ci_base-affecting files against the remote
|
||||
# image label. If hashes match the build is skipped (< 30 s); if they
|
||||
@@ -7,13 +30,16 @@ steps:
|
||||
- label: "AMD: :docker: ensure ci_base"
|
||||
key: ensure-ci-base-amd
|
||||
soft_fail: false
|
||||
depends_on: []
|
||||
depends_on:
|
||||
- refresh-rocm-base-amd
|
||||
device: amd_cpu
|
||||
no_plugin: true
|
||||
commands:
|
||||
- bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps
|
||||
- bash .buildkite/scripts/rocm/build-ci-base.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
BUILDKIT_PROGRESS: "tty"
|
||||
TERM: "xterm-256color"
|
||||
VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
|
||||
PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
|
||||
REMOTE_VLLM: "1"
|
||||
@@ -33,35 +59,12 @@ steps:
|
||||
device: amd_cpu
|
||||
no_plugin: true
|
||||
commands:
|
||||
- |
|
||||
if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" ]]; then
|
||||
echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only"
|
||||
IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts
|
||||
else
|
||||
bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel
|
||||
fi
|
||||
- |
|
||||
docker run --rm --network=none --entrypoint /bin/bash "rocm/vllm-ci:${BUILDKITE_COMMIT}" -ec '
|
||||
if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
|
||||
if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
|
||||
if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
|
||||
if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
|
||||
command -v python3
|
||||
command -v uv
|
||||
command -v pytest
|
||||
if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
|
||||
echo No ROCm CLI found in image >&2
|
||||
exit 1
|
||||
fi
|
||||
python3 - <<PY
|
||||
import torch, vllm
|
||||
print(torch.__version__)
|
||||
print(vllm.__version__)
|
||||
PY
|
||||
echo AMD image smoke OK
|
||||
'
|
||||
- bash .buildkite/scripts/rocm/build-test-image.sh
|
||||
- bash .buildkite/scripts/rocm/smoke-test-image.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
BUILDKIT_PROGRESS: "tty"
|
||||
TERM: "xterm-256color"
|
||||
VLLM_BAKE_FILE: "docker/docker-bake-rocm.hcl"
|
||||
PYTORCH_ROCM_ARCH: "gfx90a;gfx942;gfx950"
|
||||
IMAGE_TAG: "rocm/vllm-ci:$BUILDKITE_COMMIT"
|
||||
|
||||
@@ -17,12 +17,14 @@ steps:
|
||||
- tests/kernels/test_awq_int4_to_int8.py
|
||||
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
- tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
- tests/kernels/mamba/test_cpu_short_conv.py
|
||||
commands:
|
||||
- |
|
||||
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
|
||||
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
|
||||
pytest -x -v -s tests/kernels/test_onednn.py
|
||||
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
|
||||
pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
|
||||
|
||||
@@ -45,7 +45,7 @@ steps:
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
|
||||
@@ -38,7 +38,7 @@ steps:
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -76,6 +76,30 @@ steps:
|
||||
pytest -v -s v1/sample/test_logprobs.py &&
|
||||
pytest -v -s v1/sample/test_logprobs_e2e.py'
|
||||
|
||||
- label: Basic Models Tests (Initialization)
|
||||
timeout_in_minutes: 60
|
||||
device: intel_gpu
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
REPO: "vllm-ci-test-repo"
|
||||
VLLM_TEST_DEVICE: "xpu"
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/test_initialization.py
|
||||
- tests/models/registry.py
|
||||
commands:
|
||||
- >-
|
||||
bash .buildkite/scripts/hardware_ci/run-intel-test.sh
|
||||
'export VLLM_XPU_FUSED_MOE_USE_REF=1 &&
|
||||
cd tests &&
|
||||
pytest -v -s models/test_initialization.py::test_can_initialize_large_subset[Eagle3MiniMaxM2ForCausalLM]'
|
||||
|
||||
- label: XPU CPU Offload
|
||||
timeout_in_minutes: 60
|
||||
device: intel_gpu
|
||||
@@ -110,7 +134,7 @@ steps:
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 24+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
|
||||
@@ -9,7 +9,7 @@ steps:
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 2+
|
||||
mem: 24+
|
||||
mem: 16+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
|
||||
@@ -9,7 +9,7 @@ steps:
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -81,7 +81,7 @@ steps:
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
@@ -105,7 +105,7 @@ steps:
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
working_dir: "."
|
||||
env:
|
||||
|
||||
@@ -81,7 +81,7 @@ steps:
|
||||
agent_tags:
|
||||
label: production
|
||||
gpu: 1+
|
||||
mem: 16+
|
||||
mem: 24+
|
||||
no_plugin: true
|
||||
env:
|
||||
REGISTRY: "public.ecr.aws/q9t5s3a7"
|
||||
|
||||
@@ -15,7 +15,7 @@ set -euo pipefail
|
||||
|
||||
DEFAULT_REPO_SLUG="vllm-project/vllm"
|
||||
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
|
||||
DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh"
|
||||
DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
|
||||
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
|
||||
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base build_rixl build_rocshmem build_deepep mori_base ci_base"
|
||||
DEFAULT_CI_BASE_METADATA_VERSION="1"
|
||||
@@ -393,6 +393,16 @@ should_upload_wheel_artifacts() {
|
||||
|| "${TARGET}" == *"artifact"* ]]
|
||||
}
|
||||
|
||||
set_buildkite_metadata() {
|
||||
local key="$1"
|
||||
local value="$2"
|
||||
|
||||
[[ -n "${value}" ]] || return 0
|
||||
if command -v buildkite-agent >/dev/null 2>&1; then
|
||||
buildkite-agent meta-data set "${key}" "${value}" || true
|
||||
fi
|
||||
}
|
||||
|
||||
get_remote_image_label() {
|
||||
local image_ref="$1"
|
||||
local label_key="$2"
|
||||
@@ -733,6 +743,8 @@ configure_ci_base_image_refs() {
|
||||
|
||||
if is_ci_base_target; then
|
||||
IMAGE_TAG="${primary_tag}"
|
||||
CI_BASE_IMAGE="${primary_tag}"
|
||||
export CI_BASE_IMAGE
|
||||
export IMAGE_TAG
|
||||
|
||||
echo "ci_base primary image tag: ${CI_BASE_IMAGE_TAG}"
|
||||
@@ -750,6 +762,10 @@ configure_ci_base_image_refs() {
|
||||
echo "ci_base stable alias will not be pushed for this build"
|
||||
echo "Set NIGHTLY=1 on ${CI_BASE_STABLE_BRANCH:-main} to refresh ${stable_tag}"
|
||||
fi
|
||||
set_buildkite_metadata "rocm-ci-base-image" "${CI_BASE_IMAGE_TAG}"
|
||||
set_buildkite_metadata "rocm-ci-base-image-content" "${content_tag}"
|
||||
set_buildkite_metadata "rocm-ci-base-image-commit" "${CI_BASE_IMAGE_TAG_COMMIT:-}"
|
||||
set_buildkite_metadata "rocm-ci-base-image-stable" "${CI_BASE_IMAGE_TAG_STABLE:-}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
@@ -1779,7 +1795,10 @@ seed_dependency_caches_if_needed() {
|
||||
|
||||
echo "--- :docker: Seeding ${target}"
|
||||
echo "Expected cache ref: ${cache_ref}"
|
||||
docker buildx bake "${BAKE_FILES[@]}" --progress plain "${target}"
|
||||
docker buildx bake \
|
||||
"${BAKE_FILES[@]}" \
|
||||
--progress "${BUILDKIT_PROGRESS:-plain}" \
|
||||
"${target}"
|
||||
verify_dependency_cache_ref "${cache_ref}"
|
||||
done
|
||||
}
|
||||
@@ -1807,7 +1826,10 @@ run_bake() {
|
||||
local build_rc=0
|
||||
|
||||
echo "--- :docker: Building ${TARGET}"
|
||||
docker buildx bake "${BAKE_FILES[@]}" --progress plain "${BAKE_TARGETS[@]}" || build_rc=$?
|
||||
docker buildx bake \
|
||||
"${BAKE_FILES[@]}" \
|
||||
--progress "${BUILDKIT_PROGRESS:-plain}" \
|
||||
"${BAKE_TARGETS[@]}" || build_rc=$?
|
||||
|
||||
if [[ ${build_rc} -eq 0 ]]; then
|
||||
echo "--- :white_check_mark: Build complete"
|
||||
|
||||
@@ -28,6 +28,17 @@
|
||||
###############################################################################
|
||||
set -o pipefail
|
||||
|
||||
: "${BUILDKIT_PROGRESS:=plain}"
|
||||
: "${TERM:=xterm-256color}"
|
||||
: "${FORCE_COLOR:=1}"
|
||||
: "${CLICOLOR_FORCE:=1}"
|
||||
: "${PY_COLORS:=1}"
|
||||
: "${ROCM_DOCKER_TTY:=1}"
|
||||
if [[ " ${PYTEST_ADDOPTS:-} " != *" --color"* ]]; then
|
||||
PYTEST_ADDOPTS="${PYTEST_ADDOPTS:+${PYTEST_ADDOPTS} }--color=yes"
|
||||
fi
|
||||
export BUILDKIT_PROGRESS TERM FORCE_COLOR CLICOLOR_FORCE PY_COLORS PYTEST_ADDOPTS ROCM_DOCKER_TTY
|
||||
|
||||
# Export Python path for commands that run directly on the host. Containerized
|
||||
# tests set this to /vllm-workspace below so spawned Python processes do not
|
||||
# depend on their current working directory.
|
||||
@@ -149,6 +160,7 @@ EOF
|
||||
echo "--- Building local ROCm test image"
|
||||
docker build \
|
||||
--pull=false \
|
||||
--progress "${BUILDKIT_PROGRESS}" \
|
||||
--build-arg "BASE_IMAGE=${base_image}" \
|
||||
-t "${artifact_image}" \
|
||||
"${context_dir}" || return 1
|
||||
@@ -535,6 +547,13 @@ if is_multi_node "$commands"; then
|
||||
else
|
||||
echo "--- Single-node job"
|
||||
echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
|
||||
docker_run_terminal_args=(-i)
|
||||
if [[ "${ROCM_DOCKER_TTY}" == "1" ]]; then
|
||||
docker_run_terminal_args+=(-t)
|
||||
echo "Docker interactive stdin: enabled; TTY allocation: enabled"
|
||||
else
|
||||
echo "Docker interactive stdin: enabled; TTY allocation: disabled"
|
||||
fi
|
||||
|
||||
ulimit_core_hard=$(ulimit -H -c)
|
||||
if [[ "$ulimit_core_hard" == "unlimited" ]]; then
|
||||
@@ -551,6 +570,7 @@ else
|
||||
fi
|
||||
|
||||
docker run \
|
||||
"${docker_run_terminal_args[@]}" \
|
||||
--device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
|
||||
$RDMA_FLAGS \
|
||||
--network=host \
|
||||
@@ -565,6 +585,11 @@ else
|
||||
-e AWS_SECRET_ACCESS_KEY \
|
||||
-e BUILDKITE_PARALLEL_JOB \
|
||||
-e BUILDKITE_PARALLEL_JOB_COUNT \
|
||||
-e TERM \
|
||||
-e FORCE_COLOR \
|
||||
-e CLICOLOR_FORCE \
|
||||
-e PY_COLORS \
|
||||
-e PYTEST_ADDOPTS \
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
-e "HF_HOME=${HF_MOUNT}" \
|
||||
-e "PYTHONPATH=${MYPYTHONPATH}" \
|
||||
|
||||
@@ -38,7 +38,9 @@ function cpu_tests() {
|
||||
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
|
||||
pytest -x -v -s tests/kernels/core/test_cpu_activation.py
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py"
|
||||
pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
|
||||
pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py
|
||||
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py"
|
||||
|
||||
# skip tests requiring model downloads if HF_TOKEN is not set
|
||||
# due to rate-limits
|
||||
@@ -62,7 +64,6 @@ function cpu_tests() {
|
||||
set -e
|
||||
pytest -x -v -s tests/quantization/test_compressed_tensors.py::test_compressed_tensors_w8a8_logprobs"
|
||||
|
||||
|
||||
# basic online serving
|
||||
docker exec cpu-test bash -c '
|
||||
set -e
|
||||
|
||||
Executable
+32
@@ -0,0 +1,32 @@
|
||||
#!/usr/bin/env bash
|
||||
# Build the ROCm ci_base image, optionally from a freshly rebuilt ROCm base.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
metadata_get() {
|
||||
local key="$1"
|
||||
if command -v buildkite-agent >/dev/null 2>&1; then
|
||||
buildkite-agent meta-data get "${key}" 2>/dev/null || true
|
||||
fi
|
||||
}
|
||||
|
||||
main() {
|
||||
local base_refreshed=""
|
||||
|
||||
base_refreshed="$(metadata_get rocm-base-refresh)"
|
||||
if [[ "${base_refreshed}" == "1" ]]; then
|
||||
export BASE_IMAGE
|
||||
export CI_BASE_PUSH_STABLE_TAG
|
||||
|
||||
BASE_IMAGE="$(metadata_get rocm-base-image)"
|
||||
CI_BASE_PUSH_STABLE_TAG="$(metadata_get rocm-base-push-stable-tag)"
|
||||
CI_BASE_PUSH_STABLE_TAG="${CI_BASE_PUSH_STABLE_TAG:-0}"
|
||||
|
||||
echo "Using refreshed ROCm base image for ci_base: ${BASE_IMAGE}"
|
||||
echo "Push stable ci_base tag: ${CI_BASE_PUSH_STABLE_TAG}"
|
||||
fi
|
||||
|
||||
bash .buildkite/scripts/ci-bake-rocm.sh ci-base-rocm-ci-with-deps
|
||||
}
|
||||
|
||||
main "$@"
|
||||
Executable
+57
@@ -0,0 +1,57 @@
|
||||
#!/usr/bin/env bash
|
||||
# Build the ROCm CI test image or wheel artifact.
|
||||
#
|
||||
# When Dockerfile.rocm_base changes, always build the full image so downstream
|
||||
# ROCm tests can validate the freshly rebuilt base -> ci_base -> ci image chain.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
metadata_get() {
|
||||
local key="$1"
|
||||
if command -v buildkite-agent >/dev/null 2>&1; then
|
||||
buildkite-agent meta-data get "${key}" 2>/dev/null || true
|
||||
fi
|
||||
}
|
||||
|
||||
use_refreshed_base_if_present() {
|
||||
local base_refreshed=""
|
||||
|
||||
base_refreshed="$(metadata_get rocm-base-refresh)"
|
||||
if [[ "${base_refreshed}" != "1" ]]; then
|
||||
return 1
|
||||
fi
|
||||
|
||||
export BASE_IMAGE
|
||||
export CI_BASE_IMAGE
|
||||
export IMAGE_TAG_LATEST
|
||||
|
||||
BASE_IMAGE="$(metadata_get rocm-base-image)"
|
||||
CI_BASE_IMAGE="$(metadata_get rocm-ci-base-image)"
|
||||
IMAGE_TAG_LATEST="$(metadata_get rocm-ci-image-descriptive)"
|
||||
|
||||
echo "Using refreshed ROCm base image for test image: ${BASE_IMAGE}"
|
||||
echo "Using refreshed ROCm ci_base image for test image: ${CI_BASE_IMAGE}"
|
||||
if [[ -n "${IMAGE_TAG_LATEST}" ]]; then
|
||||
echo "Also tagging full ROCm CI image as: ${IMAGE_TAG_LATEST}"
|
||||
fi
|
||||
|
||||
return 0
|
||||
}
|
||||
|
||||
main() {
|
||||
local base_refreshed=0
|
||||
|
||||
if use_refreshed_base_if_present; then
|
||||
base_refreshed=1
|
||||
fi
|
||||
|
||||
if [[ "${ROCM_CI_ARTIFACT_ONLY:-0}" == "1" && "${base_refreshed}" != "1" ]]; then
|
||||
echo "ROCM_CI_ARTIFACT_ONLY=1; building ROCm wheel artifact only"
|
||||
IMAGE_TAG="" bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-artifacts
|
||||
return
|
||||
fi
|
||||
|
||||
bash .buildkite/scripts/ci-bake-rocm.sh test-rocm-ci-with-wheel
|
||||
}
|
||||
|
||||
main "$@"
|
||||
+513
@@ -0,0 +1,513 @@
|
||||
#!/usr/bin/env bash
|
||||
# Build and publish a fresh ROCm base image when Dockerfile.rocm_base changes.
|
||||
#
|
||||
# Normal AMD CI builds should not pay for this path. The script no-ops unless
|
||||
# docker/Dockerfile.rocm_base changed relative to the branch base, the previous
|
||||
# main commit, or ROCM_BASE_REFRESH_FORCE=1 is set.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
DOCKERFILE="${ROCM_BASE_DOCKERFILE:-docker/Dockerfile.rocm_base}"
|
||||
BASE_REPO="${ROCM_BASE_IMAGE_REPO:-rocm/vllm-dev}"
|
||||
CI_IMAGE_REPO="${ROCM_CI_IMAGE_REPO:-rocm/vllm-ci}"
|
||||
BUILDER_NAME="${ROCM_BASE_BUILDER_NAME:-vllm-rocm-base-builder}"
|
||||
DEFAULT_ROCM_BASE_METADATA_VERSION="1"
|
||||
DEFAULT_ROCM_BASE_CONTENT_FILES="${DOCKERFILE}"
|
||||
DEFAULT_ROCM_BASE_CONTENT_ARGS="BASE_IMAGE TRITON_BRANCH TRITON_REPO PYTORCH_BRANCH PYTORCH_REPO PYTORCH_VISION_BRANCH PYTORCH_VISION_REPO PYTORCH_AUDIO_BRANCH PYTORCH_AUDIO_REPO FA_BRANCH FA_REPO AITER_BRANCH AITER_REPO MORI_BRANCH MORI_REPO PYTORCH_ROCM_ARCH PYTHON_VERSION USE_SCCACHE"
|
||||
|
||||
metadata_set() {
|
||||
local key="$1"
|
||||
local value="$2"
|
||||
|
||||
[[ -n "${value}" ]] || return 0
|
||||
if command -v buildkite-agent >/dev/null 2>&1; then
|
||||
buildkite-agent meta-data set "${key}" "${value}" || true
|
||||
fi
|
||||
}
|
||||
|
||||
compute_content_hash() {
|
||||
local path=""
|
||||
local file=""
|
||||
|
||||
for path in "$@"; do
|
||||
if [[ -d "${path}" ]]; then
|
||||
while IFS= read -r -d '' file; do
|
||||
printf 'file:%s\n' "${file}"
|
||||
sha256sum "${file}"
|
||||
done < <(find "${path}" -type f -print0 | sort -z)
|
||||
elif [[ -f "${path}" ]]; then
|
||||
printf 'file:%s\n' "${path}"
|
||||
sha256sum "${path}"
|
||||
else
|
||||
printf 'missing:%s\n' "${path}"
|
||||
fi
|
||||
done | sha256sum | cut -d' ' -f1
|
||||
}
|
||||
|
||||
clean_docker_tag() {
|
||||
local input="$1"
|
||||
echo "${input}" | sed 's/[^a-zA-Z0-9._-]/_/g' | cut -c1-128
|
||||
}
|
||||
|
||||
tag_component() {
|
||||
local input="$1"
|
||||
local max_chars="${2:-24}"
|
||||
|
||||
clean_docker_tag "${input:-unknown}" | cut -c1-"${max_chars}"
|
||||
}
|
||||
|
||||
extract_arg_default() {
|
||||
local arg_name="$1"
|
||||
|
||||
sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
|
||||
"${DOCKERFILE}" | head -1
|
||||
}
|
||||
|
||||
resolve_image_digest() {
|
||||
local image_ref="$1"
|
||||
|
||||
docker buildx imagetools inspect "${image_ref}" 2>/dev/null \
|
||||
| sed -n -E 's/^Digest:[[:space:]]+//p' \
|
||||
| head -1 || true
|
||||
}
|
||||
|
||||
resolve_rocm_base_arg_value() {
|
||||
local arg_name="$1"
|
||||
local use_sccache="$2"
|
||||
|
||||
case "${arg_name}" in
|
||||
USE_SCCACHE)
|
||||
printf '%s\n' "${use_sccache}"
|
||||
;;
|
||||
*)
|
||||
extract_arg_default "${arg_name}"
|
||||
;;
|
||||
esac
|
||||
}
|
||||
|
||||
hash_rocm_base_arg_values() {
|
||||
local use_sccache="$1"
|
||||
local base_image_digest="$2"
|
||||
local arg_name=""
|
||||
local arg_value=""
|
||||
shift 2 || true
|
||||
|
||||
for arg_name in "$@"; do
|
||||
[[ -n "${arg_name}" ]] || continue
|
||||
arg_value=$(resolve_rocm_base_arg_value "${arg_name}" "${use_sccache}")
|
||||
printf 'arg:%s=%s\n' "${arg_name}" "${arg_value:-<empty>}"
|
||||
if [[ "${arg_name}" == "BASE_IMAGE" && -n "${arg_value}" ]]; then
|
||||
printf 'arg:%s.digest=%s\n' "${arg_name}" "${base_image_digest:-unknown}"
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
rocm_version_from_base_image() {
|
||||
local base_image="$1"
|
||||
local version=""
|
||||
|
||||
version="$(sed -n -E 's/.*:([0-9]+\.[0-9]+(\.[0-9]+)?)-.*/\1/p' <<<"${base_image}")"
|
||||
tag_component "${version:-${base_image}}" 16
|
||||
}
|
||||
|
||||
git_diff_changed_base() {
|
||||
local range="$1"
|
||||
[[ -n "$(git diff --name-only "${range}" -- "${DOCKERFILE}" 2>/dev/null)" ]]
|
||||
}
|
||||
|
||||
short_git_ref() {
|
||||
local ref="$1"
|
||||
|
||||
git rev-parse --short "${ref}" 2>/dev/null || printf '%s\n' "${ref}"
|
||||
}
|
||||
|
||||
extract_arg_default_from_ref() {
|
||||
local ref="$1"
|
||||
local arg_name="$2"
|
||||
local content=""
|
||||
|
||||
content="$(git show "${ref}:${DOCKERFILE}" 2>/dev/null || true)"
|
||||
sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \
|
||||
<<<"${content}" | head -1
|
||||
}
|
||||
|
||||
log_arg_default_changes() {
|
||||
local old_ref="$1"
|
||||
local new_ref="$2"
|
||||
local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}"
|
||||
local arg_name=""
|
||||
local old_value=""
|
||||
local new_value=""
|
||||
local changed=0
|
||||
|
||||
echo "Changed ROCm base ARG defaults:"
|
||||
for arg_name in ${content_args}; do
|
||||
old_value="$(extract_arg_default_from_ref "${old_ref}" "${arg_name}")"
|
||||
new_value="$(extract_arg_default_from_ref "${new_ref}" "${arg_name}")"
|
||||
if [[ "${old_value}" != "${new_value}" ]]; then
|
||||
echo " - ${arg_name}: ${old_value:-<unset>} -> ${new_value:-<unset>}"
|
||||
changed=1
|
||||
fi
|
||||
done
|
||||
|
||||
if [[ "${changed}" == "0" ]]; then
|
||||
echo " - none detected; Dockerfile instructions changed outside tracked ARG defaults"
|
||||
fi
|
||||
}
|
||||
|
||||
log_arg_line_diff() {
|
||||
local range="$1"
|
||||
local arg_diff=""
|
||||
|
||||
arg_diff="$(
|
||||
git diff --unified=0 "${range}" -- "${DOCKERFILE}" 2>/dev/null \
|
||||
| awk '/^[+-][[:space:]]*ARG[[:space:]]/ && $0 !~ /^(---|\+\+\+)/ { print " " $0 }' \
|
||||
|| true
|
||||
)"
|
||||
|
||||
if [[ -n "${arg_diff}" ]]; then
|
||||
echo "Changed Dockerfile ARG lines:"
|
||||
printf '%s\n' "${arg_diff}"
|
||||
fi
|
||||
}
|
||||
|
||||
log_rocm_base_change_check() {
|
||||
local context="$1"
|
||||
local range="$2"
|
||||
local old_ref="$3"
|
||||
local old_short=""
|
||||
local head_short=""
|
||||
|
||||
old_short="$(short_git_ref "${old_ref}")"
|
||||
head_short="$(short_git_ref HEAD)"
|
||||
|
||||
echo "--- :mag: ROCm base refresh check"
|
||||
echo "Context: ${context}"
|
||||
echo "Dockerfile: ${DOCKERFILE}"
|
||||
echo "Base revision: ${old_short}"
|
||||
echo "Head revision: ${head_short}"
|
||||
echo "Git diff range: ${range}"
|
||||
}
|
||||
|
||||
log_rocm_base_rebuild_reason() {
|
||||
local context="$1"
|
||||
local range="$2"
|
||||
local old_ref="$3"
|
||||
local changed_files=""
|
||||
|
||||
log_rocm_base_change_check "${context}" "${range}" "${old_ref}"
|
||||
|
||||
changed_files="$(git diff --name-only "${range}" -- "${DOCKERFILE}" 2>/dev/null || true)"
|
||||
echo "Changed files:"
|
||||
if [[ -n "${changed_files}" ]]; then
|
||||
sed 's/^/ - /' <<<"${changed_files}"
|
||||
else
|
||||
echo " - ${DOCKERFILE}"
|
||||
fi
|
||||
log_arg_default_changes "${old_ref}" HEAD
|
||||
log_arg_line_diff "${range}"
|
||||
echo "Decision: rebuilding ROCm base image because ${DOCKERFILE} changed."
|
||||
}
|
||||
|
||||
rocm_base_changed_in_range() {
|
||||
local context="$1"
|
||||
local range="$2"
|
||||
local old_ref="$3"
|
||||
|
||||
if git_diff_changed_base "${range}"; then
|
||||
log_rocm_base_rebuild_reason "${context}" "${range}" "${old_ref}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
log_rocm_base_change_check "${context}" "${range}" "${old_ref}"
|
||||
echo "Decision: ROCm base refresh not required; ${DOCKERFILE} is unchanged."
|
||||
return 1
|
||||
}
|
||||
|
||||
rocm_base_changed() {
|
||||
local base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
|
||||
local base_ref="refs/remotes/origin/${base_branch}"
|
||||
local merge_base=""
|
||||
|
||||
if [[ "${ROCM_BASE_REFRESH_SKIP:-0}" == "1" ]]; then
|
||||
echo "ROCM_BASE_REFRESH_SKIP=1 set; skipping ROCm base refresh"
|
||||
return 1
|
||||
fi
|
||||
|
||||
if [[ "${ROCM_BASE_REFRESH_FORCE:-0}" == "1" ]]; then
|
||||
echo "ROCM_BASE_REFRESH_FORCE=1 set; refreshing ROCm base image"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if ! git rev-parse --is-inside-work-tree >/dev/null 2>&1; then
|
||||
echo "Not in a git checkout; skipping ROCm base refresh unless forced"
|
||||
return 1
|
||||
fi
|
||||
|
||||
if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
|
||||
git fetch --no-tags --depth=200 origin \
|
||||
"+refs/heads/${base_branch}:${base_ref}" >/dev/null 2>&1 || true
|
||||
merge_base=$(git merge-base HEAD "${base_ref}" 2>/dev/null || true)
|
||||
if [[ -z "${merge_base}" ]]; then
|
||||
echo "Unable to determine merge base with PR base ${base_ref}; skipping ROCm base refresh unless forced"
|
||||
return 1
|
||||
fi
|
||||
if rocm_base_changed_in_range \
|
||||
"pull request build against ${base_ref}" \
|
||||
"${merge_base}...HEAD" \
|
||||
"${merge_base}"; then
|
||||
return 0
|
||||
fi
|
||||
elif [[ "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]] \
|
||||
&& git rev-parse --verify HEAD~1 >/dev/null 2>&1; then
|
||||
if rocm_base_changed_in_range \
|
||||
"stable branch build; comparing against previous ${ROCM_BASE_STABLE_BRANCH:-main} commit" \
|
||||
"HEAD~1..HEAD" \
|
||||
"HEAD~1"; then
|
||||
return 0
|
||||
fi
|
||||
else
|
||||
git fetch --no-tags --depth=200 origin \
|
||||
"+refs/heads/${base_branch}:${base_ref}" >/dev/null 2>&1 || true
|
||||
merge_base=$(git merge-base HEAD "${base_ref}" 2>/dev/null || true)
|
||||
if [[ -z "${merge_base}" ]]; then
|
||||
echo "Unable to determine merge base with branch base ${base_ref}; skipping ROCm base refresh unless forced"
|
||||
return 1
|
||||
fi
|
||||
if rocm_base_changed_in_range \
|
||||
"branch build against ${base_ref}" \
|
||||
"${merge_base}...HEAD" \
|
||||
"${merge_base}"; then
|
||||
return 0
|
||||
fi
|
||||
fi
|
||||
|
||||
return 1
|
||||
}
|
||||
|
||||
should_push_stable_tag() {
|
||||
if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" ]]; then
|
||||
return 1
|
||||
fi
|
||||
|
||||
if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "1" ]]; then
|
||||
return 0
|
||||
fi
|
||||
if [[ "${ROCM_BASE_PUSH_STABLE_TAG:-}" == "0" ]]; then
|
||||
return 1
|
||||
fi
|
||||
|
||||
[[ "${BUILDKITE_PULL_REQUEST:-false}" == "false" \
|
||||
&& "${BUILDKITE_BRANCH:-}" == "${ROCM_BASE_STABLE_BRANCH:-main}" ]]
|
||||
}
|
||||
|
||||
setup_builder() {
|
||||
echo "--- :buildkite: Setting up buildx builder for ROCm base"
|
||||
if docker buildx inspect "${BUILDER_NAME}" >/dev/null 2>&1; then
|
||||
docker buildx use "${BUILDER_NAME}"
|
||||
else
|
||||
docker buildx create --name "${BUILDER_NAME}" --driver docker-container --use
|
||||
fi
|
||||
docker buildx inspect --bootstrap
|
||||
}
|
||||
|
||||
compute_base_content_hash() {
|
||||
local use_sccache="$1"
|
||||
local base_image_digest="$2"
|
||||
local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}"
|
||||
local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}"
|
||||
local -a content_paths=()
|
||||
local -a content_arg_names=()
|
||||
|
||||
read -r -a content_paths <<< "${content_files}"
|
||||
read -r -a content_arg_names <<< "${content_args}"
|
||||
|
||||
{
|
||||
printf 'content-files-hash:%s\n' "$(compute_content_hash "${content_paths[@]}")"
|
||||
printf 'dockerfile:%s\n' "${DOCKERFILE}"
|
||||
printf 'resolved-build-args:\n'
|
||||
hash_rocm_base_arg_values \
|
||||
"${use_sccache}" "${base_image_digest}" "${content_arg_names[@]}"
|
||||
} | sha256sum | cut -d' ' -f1
|
||||
}
|
||||
|
||||
build_base_image() {
|
||||
local use_sccache="${ROCM_BASE_USE_SCCACHE:-${USE_SCCACHE:-0}}"
|
||||
local base_hash=""
|
||||
local build_date=""
|
||||
local build_suffix=""
|
||||
local base_image_arg=""
|
||||
local base_image_digest=""
|
||||
local rocm_version=""
|
||||
local triton_arg=""
|
||||
local pytorch_arg=""
|
||||
local pytorch_vision_arg=""
|
||||
local pytorch_audio_arg=""
|
||||
local fa_arg=""
|
||||
local aiter_arg=""
|
||||
local mori_arg=""
|
||||
local python_version_arg=""
|
||||
local pytorch_rocm_arch_arg=""
|
||||
local pytorch_branch=""
|
||||
local aiter_branch=""
|
||||
local dependency_summary=""
|
||||
local descriptor=""
|
||||
local ci_descriptor=""
|
||||
local descriptive_tag=""
|
||||
local stable_tag="${BASE_REPO}:base"
|
||||
local ci_descriptive_tag=""
|
||||
local content_files="${ROCM_BASE_CONTENT_FILES:-${DEFAULT_ROCM_BASE_CONTENT_FILES}}"
|
||||
local content_args="${ROCM_BASE_CONTENT_ARGS:-${DEFAULT_ROCM_BASE_CONTENT_ARGS}}"
|
||||
local content_files_hash=""
|
||||
local metadata_version="${ROCM_BASE_METADATA_VERSION:-${DEFAULT_ROCM_BASE_METADATA_VERSION}}"
|
||||
local -a tags=()
|
||||
local -a no_cache_args=()
|
||||
local -a sccache_args=()
|
||||
local -a content_paths=()
|
||||
|
||||
if [[ ! -f "${DOCKERFILE}" ]]; then
|
||||
echo "Error: ROCm base Dockerfile not found: ${DOCKERFILE}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
build_date="${ROCM_BASE_TAG_DATE:-$(date -u +%Y%m%d)}"
|
||||
if [[ -n "${BUILDKITE_BUILD_NUMBER:-}" ]]; then
|
||||
build_suffix="_bk_${BUILDKITE_BUILD_NUMBER}"
|
||||
fi
|
||||
base_image_arg="$(extract_arg_default BASE_IMAGE)"
|
||||
base_image_digest="$(resolve_image_digest "${base_image_arg}")"
|
||||
read -r -a content_paths <<< "${content_files}"
|
||||
content_files_hash="$(compute_content_hash "${content_paths[@]}")"
|
||||
base_hash=$(compute_base_content_hash "${use_sccache}" "${base_image_digest}")
|
||||
rocm_version="$(rocm_version_from_base_image "${base_image_arg}")"
|
||||
triton_arg="$(extract_arg_default TRITON_BRANCH)"
|
||||
pytorch_arg="$(extract_arg_default PYTORCH_BRANCH)"
|
||||
pytorch_vision_arg="$(extract_arg_default PYTORCH_VISION_BRANCH)"
|
||||
pytorch_audio_arg="$(extract_arg_default PYTORCH_AUDIO_BRANCH)"
|
||||
fa_arg="$(extract_arg_default FA_BRANCH)"
|
||||
aiter_arg="$(extract_arg_default AITER_BRANCH)"
|
||||
mori_arg="$(extract_arg_default MORI_BRANCH)"
|
||||
python_version_arg="$(extract_arg_default PYTHON_VERSION)"
|
||||
pytorch_rocm_arch_arg="$(extract_arg_default PYTORCH_ROCM_ARCH)"
|
||||
pytorch_branch="$(tag_component "${pytorch_arg}" 16)"
|
||||
aiter_branch="$(tag_component "${aiter_arg}" 24)"
|
||||
dependency_summary="base=${base_image_arg},rocm=${rocm_version},python=${python_version_arg},pytorch=${pytorch_arg},torchvision=${pytorch_vision_arg},torchaudio=${pytorch_audio_arg},triton=${triton_arg},flash-attn=${fa_arg},aiter=${aiter_arg},mori=${mori_arg},pytorch-rocm-arch=${pytorch_rocm_arch_arg}"
|
||||
descriptor="$(clean_docker_tag "base_custom_aiter_${aiter_branch}_torch_${pytorch_branch}_${build_date}${build_suffix}")"
|
||||
ci_descriptor="$(clean_docker_tag "ci_custom_aiter_${aiter_branch}_torch_${pytorch_branch}_${build_date}${build_suffix}")"
|
||||
|
||||
descriptive_tag="${BASE_REPO}:${descriptor}"
|
||||
ci_descriptive_tag="${CI_IMAGE_REPO}:${ci_descriptor}"
|
||||
|
||||
tags=(-t "${descriptive_tag}")
|
||||
if should_push_stable_tag; then
|
||||
tags+=(-t "${stable_tag}")
|
||||
metadata_set "rocm-base-push-stable-tag" "1"
|
||||
else
|
||||
metadata_set "rocm-base-push-stable-tag" "0"
|
||||
fi
|
||||
|
||||
if [[ "${ROCM_BASE_NO_CACHE:-1}" == "1" ]]; then
|
||||
no_cache_args=(--no-cache)
|
||||
fi
|
||||
|
||||
for env_name in \
|
||||
SCCACHE_DOWNLOAD_URL \
|
||||
SCCACHE_ENDPOINT \
|
||||
SCCACHE_BUCKET_NAME \
|
||||
SCCACHE_REGION_NAME \
|
||||
SCCACHE_S3_NO_CREDENTIALS; do
|
||||
if [[ -n "${!env_name:-}" ]]; then
|
||||
sccache_args+=(--build-arg "${env_name}=${!env_name}")
|
||||
fi
|
||||
done
|
||||
|
||||
echo "--- :docker: Building ROCm base image"
|
||||
echo "Dockerfile: ${DOCKERFILE}"
|
||||
echo "Descriptive tag: ${descriptive_tag}"
|
||||
echo "Stable tag: ${stable_tag} ($(should_push_stable_tag && echo enabled || echo disabled))"
|
||||
echo "Content hash: ${base_hash}"
|
||||
echo "Dependency summary: ${dependency_summary}"
|
||||
echo "USE_SCCACHE: ${use_sccache}"
|
||||
|
||||
docker buildx build \
|
||||
"${no_cache_args[@]}" \
|
||||
--pull \
|
||||
--progress "${BUILDKIT_PROGRESS:-plain}" \
|
||||
--file "${DOCKERFILE}" \
|
||||
--build-arg "USE_SCCACHE=${use_sccache}" \
|
||||
"${sccache_args[@]}" \
|
||||
--label "org.opencontainers.image.source=https://github.com/vllm-project/vllm" \
|
||||
--label "org.opencontainers.image.vendor=vLLM" \
|
||||
--label "org.opencontainers.image.title=vLLM ROCm base" \
|
||||
--label "org.opencontainers.image.revision=${BUILDKITE_COMMIT:-}" \
|
||||
--label "vllm.rocm_base.metadata_version=${metadata_version}" \
|
||||
--label "vllm.rocm_base.content_hash=${base_hash}" \
|
||||
--label "vllm.rocm_base.content_files_hash=${content_files_hash}" \
|
||||
--label "vllm.rocm_base.dockerfile=${DOCKERFILE}" \
|
||||
--label "vllm.rocm_base.image.descriptive=${descriptive_tag}" \
|
||||
--label "vllm.rocm_base.image.stable=${stable_tag}" \
|
||||
--label "vllm.rocm_base.git_commit=${BUILDKITE_COMMIT:-}" \
|
||||
--label "vllm.rocm_base.stable_branch=${ROCM_BASE_STABLE_BRANCH:-main}" \
|
||||
--label "vllm.rocm_base.descriptor=${descriptor}" \
|
||||
--label "vllm.rocm_base.dependency_summary=${dependency_summary}" \
|
||||
--label "vllm.rocm_base.base_image=${base_image_arg}" \
|
||||
--label "vllm.rocm_base.base_image_digest=${base_image_digest}" \
|
||||
--label "vllm.rocm_base.dependency.rocm=${rocm_version}" \
|
||||
--label "vllm.rocm_base.dependency.python=${python_version_arg}" \
|
||||
--label "vllm.rocm_base.dependency.pytorch=${pytorch_arg}" \
|
||||
--label "vllm.rocm_base.dependency.torchvision=${pytorch_vision_arg}" \
|
||||
--label "vllm.rocm_base.dependency.torchaudio=${pytorch_audio_arg}" \
|
||||
--label "vllm.rocm_base.dependency.triton=${triton_arg}" \
|
||||
--label "vllm.rocm_base.dependency.flash_attention=${fa_arg}" \
|
||||
--label "vllm.rocm_base.dependency.aiter=${aiter_arg}" \
|
||||
--label "vllm.rocm_base.dependency.mori=${mori_arg}" \
|
||||
--label "vllm.rocm_base.pytorch_rocm_arch=${pytorch_rocm_arch_arg}" \
|
||||
"${tags[@]}" \
|
||||
--push \
|
||||
.
|
||||
|
||||
docker buildx imagetools inspect "${descriptive_tag}" >/dev/null
|
||||
|
||||
metadata_set "rocm-base-refresh" "1"
|
||||
metadata_set "rocm-base-image" "${descriptive_tag}"
|
||||
metadata_set "rocm-base-image-descriptive" "${descriptive_tag}"
|
||||
metadata_set "rocm-base-image-stable" "${stable_tag}"
|
||||
metadata_set "rocm-base-image-ci-descriptive" "${ci_descriptive_tag}"
|
||||
metadata_set "rocm-base-metadata-version" "${metadata_version}"
|
||||
metadata_set "rocm-base-content-hash" "${base_hash}"
|
||||
metadata_set "rocm-base-content-files-hash" "${content_files_hash}"
|
||||
metadata_set "rocm-base-content-files" "${content_files}"
|
||||
metadata_set "rocm-base-content-args" "${content_args}"
|
||||
metadata_set "rocm-base-base-image-digest" "${base_image_digest}"
|
||||
metadata_set "rocm-base-dockerfile" "${DOCKERFILE}"
|
||||
metadata_set "rocm-base-descriptor" "${descriptor}"
|
||||
metadata_set "rocm-base-dependency-summary" "${dependency_summary}"
|
||||
metadata_set "rocm-base-dependency-rocm" "${rocm_version}"
|
||||
metadata_set "rocm-base-dependency-python" "${python_version_arg}"
|
||||
metadata_set "rocm-base-dependency-pytorch" "${pytorch_arg}"
|
||||
metadata_set "rocm-base-dependency-torchvision" "${pytorch_vision_arg}"
|
||||
metadata_set "rocm-base-dependency-torchaudio" "${pytorch_audio_arg}"
|
||||
metadata_set "rocm-base-dependency-triton" "${triton_arg}"
|
||||
metadata_set "rocm-base-dependency-flash-attention" "${fa_arg}"
|
||||
metadata_set "rocm-base-dependency-aiter" "${aiter_arg}"
|
||||
metadata_set "rocm-base-dependency-mori" "${mori_arg}"
|
||||
metadata_set "rocm-base-pytorch-rocm-arch" "${pytorch_rocm_arch_arg}"
|
||||
metadata_set "rocm-ci-image-descriptive" "${ci_descriptive_tag}"
|
||||
|
||||
echo "--- :white_check_mark: ROCm base image published"
|
||||
echo "Use BASE_IMAGE=${descriptive_tag} for downstream ROCm CI builds"
|
||||
}
|
||||
|
||||
main() {
|
||||
metadata_set "rocm-base-refresh" "0"
|
||||
|
||||
if ! rocm_base_changed; then
|
||||
echo "ROCm base Dockerfile did not change; skipping base image refresh"
|
||||
return 0
|
||||
fi
|
||||
|
||||
setup_builder
|
||||
build_base_image
|
||||
}
|
||||
|
||||
main "$@"
|
||||
Executable
+32
@@ -0,0 +1,32 @@
|
||||
#!/usr/bin/env bash
|
||||
# Fast structural smoke test for the full ROCm CI image.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
image_ref="${VLLM_CI_SMOKE_IMAGE:-rocm/vllm-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}}"
|
||||
|
||||
docker run --rm --network=none --entrypoint /bin/bash "${image_ref}" -ec '
|
||||
if [ ! -d /vllm-workspace ]; then echo Missing directory: /vllm-workspace >&2; exit 1; fi
|
||||
if [ ! -d /vllm-workspace/tests ]; then echo Missing directory: /vllm-workspace/tests >&2; exit 1; fi
|
||||
if [ ! -d /vllm-workspace/src/vllm ]; then echo Missing directory: /vllm-workspace/src/vllm >&2; exit 1; fi
|
||||
if [ ! -x /vllm-workspace/src/vllm/vllm-rs ]; then echo Missing executable: /vllm-workspace/src/vllm/vllm-rs >&2; exit 1; fi
|
||||
|
||||
command -v python3
|
||||
command -v uv
|
||||
command -v pytest
|
||||
|
||||
if ! command -v amd-smi >/dev/null 2>&1 && ! command -v rocminfo >/dev/null 2>&1; then
|
||||
echo No ROCm CLI found in image >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
python3 - <<PY
|
||||
import torch
|
||||
import vllm
|
||||
|
||||
print(torch.__version__)
|
||||
print(vllm.__version__)
|
||||
PY
|
||||
|
||||
echo AMD image smoke OK
|
||||
'
|
||||
@@ -109,7 +109,9 @@ run_nodes() {
|
||||
if [ "$node" -ne 0 ]; then
|
||||
docker exec -d "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
|
||||
else
|
||||
docker exec "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
|
||||
# Allocate a TTY (-t -i) for the foreground head node so its output
|
||||
# keeps ANSI color in the Buildkite log (see run-amd-test.sh).
|
||||
docker exec -t -i "node$node" /bin/bash -c "cd $WORKING_DIR ; ${COMMANDS[$node]}"
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
@@ -8,7 +8,12 @@ if [[ "$MODE" != "style-clippy" && "$MODE" != "test" ]]; then
|
||||
exit 2
|
||||
fi
|
||||
|
||||
ROOT_DIR="$(git rev-parse --show-toplevel)"
|
||||
if ROOT_DIR="$(git rev-parse --show-toplevel 2>/dev/null)"; then
|
||||
:
|
||||
else
|
||||
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
|
||||
ROOT_DIR="$(cd -- "${SCRIPT_DIR}/../.." && pwd -P)"
|
||||
fi
|
||||
cd "$ROOT_DIR"
|
||||
|
||||
export CARGO_TERM_COLOR="${CARGO_TERM_COLOR:-always}"
|
||||
|
||||
+222
-11
@@ -380,6 +380,7 @@ steps:
|
||||
- tests/test_outputs.py
|
||||
- tests/test_pooling_params.py
|
||||
- tests/test_ray_env.py
|
||||
- tests/test_sampling_params.py
|
||||
- tests/multimodal
|
||||
- tests/renderers
|
||||
- tests/standalone_tests/lazy_imports.py
|
||||
@@ -395,6 +396,7 @@ steps:
|
||||
- pytest -v -s test_outputs.py
|
||||
- pytest -v -s test_pooling_params.py
|
||||
- pytest -v -s test_ray_env.py
|
||||
- pytest -v -s test_sampling_params.py
|
||||
- pytest -v -s -m 'cpu_test' multimodal
|
||||
- pytest -v -s renderers
|
||||
- pytest -v -s tokenizers_
|
||||
@@ -404,6 +406,36 @@ steps:
|
||||
- pytest -v -s transformers_utils
|
||||
- pytest -v -s config
|
||||
|
||||
#------------------------------------------------------------ mi250 · rust -----------------------------------------------------------#
|
||||
|
||||
- label: Rust Frontend Cargo Style + Clippy # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
no_gpu: true
|
||||
working_dir: "/vllm-workspace"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/test_areas/rust_frontend_cargo.yaml
|
||||
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh
|
||||
commands:
|
||||
- bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh style-clippy
|
||||
|
||||
- label: Rust Frontend Cargo Tests # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
|
||||
agent_pool: mi250_1
|
||||
no_gpu: true
|
||||
working_dir: "/vllm-workspace"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- rust-toolchain.toml
|
||||
- .buildkite/test_areas/rust_frontend_cargo.yaml
|
||||
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh
|
||||
commands:
|
||||
- bash .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
|
||||
|
||||
#----------------------------------------------------------- mi250 · docker ----------------------------------------------------------#
|
||||
|
||||
- label: Docker Build Metadata (ROCm) # TBD
|
||||
@@ -472,7 +504,7 @@ steps:
|
||||
commands:
|
||||
- TARGET_TEST_SUITE=MI300 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
|
||||
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
|
||||
- pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
|
||||
- pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
|
||||
@@ -1301,7 +1333,7 @@ steps:
|
||||
#---------------------------------------------------------- mi300 · kernels ----------------------------------------------------------#
|
||||
|
||||
- label: vLLM IR Tests # TBD
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1350,7 +1382,7 @@ steps:
|
||||
- pytest -v -s kernels/core --ignore=kernels/core/test_minimax_reduce_rms.py kernels/test_concat_mla_q.py kernels/test_top_k_per_row.py
|
||||
|
||||
- label: Kernels KDA Test # TBD
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1558,7 +1590,7 @@ steps:
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_eagle_dp.py
|
||||
|
||||
- label: Model Runner V2 Pipeline Parallelism (4 GPUs) # TBD
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_4
|
||||
num_gpus: 4
|
||||
@@ -1577,7 +1609,7 @@ steps:
|
||||
- pytest -v -s distributed/test_pp_cudagraph.py -k "not ray"
|
||||
|
||||
- label: Model Runner V2 Spec Decode # TBD
|
||||
timeout_in_minutes: 45
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1604,7 +1636,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
parallelism: 2
|
||||
parallelism: 6
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/models/
|
||||
@@ -1637,10 +1669,10 @@ steps:
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/test_terratorch.py
|
||||
- tests/models/test_transformers.py
|
||||
- tests/models/transformers/test_backend.py
|
||||
- tests/models/test_registry.py
|
||||
commands:
|
||||
- pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
|
||||
- pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
|
||||
|
||||
#----------------------------------------------------- mi300 · models / language -----------------------------------------------------#
|
||||
|
||||
@@ -1860,7 +1892,7 @@ steps:
|
||||
- examples/
|
||||
commands:
|
||||
- pip install --upgrade git+https://github.com/huggingface/transformers
|
||||
- pytest -v -s tests/models/test_transformers.py
|
||||
- pytest -v -s tests/models/transformers/test_backend.py
|
||||
- pytest -v -s tests/models/multimodal/test_mapping.py
|
||||
- python3 examples/basic/offline_inference/chat.py
|
||||
- python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl
|
||||
@@ -1915,7 +1947,7 @@ steps:
|
||||
- pytest -v -s plugins_tests/lora_resolvers # unit tests for in-tree lora resolver plugins
|
||||
|
||||
- label: GGUF Plugin # TBD
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
@@ -1930,6 +1962,128 @@ steps:
|
||||
- pip install "vllm-gguf-plugin >= 0.0.2"
|
||||
- pytest -v -s plugins_tests/gguf
|
||||
|
||||
#------------------------------------------------------- mi300 · rust_frontend -------------------------------------------------------#
|
||||
|
||||
- label: Rust Frontend OpenAI Coverage # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- vllm/benchmarks/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/entrypoints/serve/
|
||||
- vllm/v1/sample/
|
||||
- tests/utils.py
|
||||
- tests/benchmarks/test_serve_cli.py
|
||||
- tests/entrypoints/openai/chat_completion/test_chat_completion.py
|
||||
- tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
|
||||
- tests/entrypoints/openai/completion/test_shutdown.py
|
||||
- tests/entrypoints/openai/test_return_token_ids.py
|
||||
- tests/entrypoints/openai/test_uds.py
|
||||
- tests/v1/sample/test_logprobs_e2e.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
|
||||
- pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
|
||||
- pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison"
|
||||
- pytest -v -s entrypoints/openai/test_uds.py
|
||||
- pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
|
||||
|
||||
- label: Rust Frontend Serve Admin Coverage # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/entrypoints/serve/
|
||||
- vllm/v1/engine/
|
||||
- tests/utils.py
|
||||
- tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
- tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
|
||||
- tests/entrypoints/serve/instrumentator/test_basic.py
|
||||
- tests/entrypoints/serve/instrumentator/test_metrics.py
|
||||
- tests/entrypoints/serve/tokenize/test_tokenization.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
|
||||
- pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
|
||||
- pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info and not test_tokenize_chat"
|
||||
|
||||
- label: Rust Frontend Core Correctness # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- vllm/entrypoints/openai/
|
||||
- tests/utils.py
|
||||
- tests/entrypoints/openai/correctness/test_lmeval.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s entrypoints/openai/correctness/test_lmeval.py::test_lm_eval_accuracy_v1_engine
|
||||
|
||||
- label: Rust Frontend Tool Use # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_1
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- vllm/entrypoints/openai/
|
||||
- vllm/tool_parsers/
|
||||
- tests/utils.py
|
||||
- tests/tool_use/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s tool_use --ignore=tool_use/mistral --models llama3.2 -k "not test_response_format_with_tool_choice_required and not test_parallel_tool_calls_false and not test_tool_call_and_choice"
|
||||
|
||||
- label: Rust Frontend Distributed # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_4
|
||||
num_gpus: 4
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- rust/
|
||||
- vllm/distributed/
|
||||
- vllm/engine/
|
||||
- vllm/executor/
|
||||
- vllm/v1/engine/
|
||||
- vllm/v1/worker/
|
||||
- tests/utils.py
|
||||
- tests/v1/distributed/test_external_lb_dp.py
|
||||
- tests/v1/distributed/test_hybrid_lb_dp.py
|
||||
- tests/v1/distributed/test_internal_lb_dp.py
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_internal_lb_dp.py -k "not 4 and not server_info"
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/distributed/test_external_lb_dp.py -k "not 4 and not server_info"
|
||||
- TP_SIZE=1 DP_SIZE=4 pytest -v -s v1/distributed/test_hybrid_lb_dp.py -k "not 4 and not server_info"
|
||||
|
||||
#------------------------------------------------------- mi300 · quantization --------------------------------------------------------#
|
||||
|
||||
- label: Quantization # TBD
|
||||
@@ -2439,6 +2593,59 @@ steps:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
|
||||
|
||||
- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- vllm/v1/core/sched/
|
||||
- vllm/v1/core/kv_cache_coordinator.py
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
|
||||
|
||||
- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
agent_pool: mi300_2
|
||||
num_gpus: 2
|
||||
optional: true
|
||||
working_dir: "/vllm-workspace/tests"
|
||||
source_file_dependencies:
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/offloading/
|
||||
- tests/v1/kv_connector/nixl_integration/
|
||||
- vllm/platforms/rocm.py
|
||||
commands:
|
||||
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
|
||||
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
|
||||
|
||||
- label: V1 e2e (4 GPUs) # TBD
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
|
||||
@@ -2729,15 +2936,19 @@ steps:
|
||||
- vllm/envs.py
|
||||
- examples/offline_inference/data_parallel.py
|
||||
- tests/distributed/test_context_parallel.py
|
||||
- tests/distributed/test_rocm_aiter_custom_ar.py
|
||||
- tests/distributed/test_rocm_quick_reduce.py
|
||||
- tests/distributed/test_quick_all_reduce.py
|
||||
- tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
|
||||
- tests/v1/distributed/test_dbo.py
|
||||
- tests/utils.py
|
||||
commands:
|
||||
- pytest -v -s tests/distributed/test_context_parallel.py
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
- pytest -v -s tests/distributed/test_rocm_aiter_custom_ar.py
|
||||
- pytest -v -s tests/v1/e2e/general/test_rocm_aiter_custom_ar.py
|
||||
- pytest -v -s tests/distributed/test_rocm_quick_reduce.py
|
||||
- pytest -v -s tests/distributed/test_quick_all_reduce.py
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
|
||||
#-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------#
|
||||
|
||||
|
||||
@@ -150,9 +150,9 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
|
||||
|
||||
- label: LM Eval Humming (A100 - TEMPORARY)
|
||||
key: lm-eval-humming-a100
|
||||
timeout_in_minutes: 30
|
||||
- label: LM Eval Humming f16 (A100 - TEMPORARY)
|
||||
key: lm-eval-humming-f16-a100
|
||||
timeout_in_minutes: 120
|
||||
device: a100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
@@ -160,13 +160,29 @@ steps:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
|
||||
|
||||
- label: LM Eval Humming (H100 - TEMPORARY)
|
||||
key: lm-eval-humming-h100
|
||||
timeout_in_minutes: 30
|
||||
- label: LM Eval Humming Act int8 (A100 - TEMPORARY)
|
||||
key: lm-eval-humming-act-a100
|
||||
timeout_in_minutes: 120
|
||||
device: a100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
|
||||
|
||||
- label: LM Eval Humming f16 (H100 - TEMPORARY)
|
||||
key: lm-eval-humming-f16-h100
|
||||
timeout_in_minutes: 120
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
@@ -174,14 +190,30 @@ steps:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
|
||||
|
||||
- label: LM Eval Humming (B200 - TEMPORARY)
|
||||
key: lm-eval-humming-b200
|
||||
timeout_in_minutes: 30
|
||||
- label: LM Eval Humming Act fp8/int8 (H100 - TEMPORARY)
|
||||
key: lm-eval-humming-act-h100
|
||||
timeout_in_minutes: 120
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
|
||||
|
||||
- label: LM Eval Humming f16 (B200 - TEMPORARY)
|
||||
key: lm-eval-humming-f16-b200
|
||||
timeout_in_minutes: 120
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 1
|
||||
@@ -189,10 +221,26 @@ steps:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config.txt
|
||||
|
||||
- label: LM Eval Humming Act fp8/int8 (B200 - TEMPORARY)
|
||||
key: lm-eval-humming-act-b200
|
||||
timeout_in_minutes: 120
|
||||
device: b200-k8s
|
||||
optional: true
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- vllm/model_executor/layers/quantization/humming.py
|
||||
- vllm/model_executor/layers/quantization/utils/humming_utils.py
|
||||
- vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
|
||||
- vllm/model_executor/layers/fused_moe/oracle/
|
||||
- vllm/model_executor/kernels/linear/
|
||||
commands:
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-fp8.txt
|
||||
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/humming/config-act-int8.txt
|
||||
|
||||
- label: LM Eval TurboQuant KV Cache
|
||||
key: lm-eval-turboquant-kv-cache
|
||||
|
||||
@@ -351,6 +351,7 @@ steps:
|
||||
- tests/test_outputs.py
|
||||
- tests/test_pooling_params.py
|
||||
- tests/test_ray_env.py
|
||||
- tests/test_sampling_params.py
|
||||
- tests/multimodal
|
||||
- tests/renderers
|
||||
- tests/standalone_tests/lazy_imports.py
|
||||
@@ -368,6 +369,7 @@ steps:
|
||||
- pytest -v -s test_outputs.py
|
||||
- pytest -v -s test_pooling_params.py
|
||||
- pytest -v -s test_ray_env.py
|
||||
- pytest -v -s test_sampling_params.py
|
||||
- pytest -v -s -m 'cpu_test' multimodal
|
||||
- pytest -v -s renderers
|
||||
- pytest -v -s reasoning
|
||||
|
||||
@@ -36,10 +36,10 @@ steps:
|
||||
source_file_dependencies:
|
||||
- vllm/
|
||||
- tests/models/test_terratorch.py
|
||||
- tests/models/test_transformers.py
|
||||
- tests/models/transformers/test_backend.py
|
||||
- tests/models/test_registry.py
|
||||
commands:
|
||||
- pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
|
||||
- pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py
|
||||
mirror:
|
||||
amd:
|
||||
device: mi325_1
|
||||
@@ -55,6 +55,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/test_utils.py
|
||||
- tests/models/test_vision.py
|
||||
- tests/models/transformers/fusers/
|
||||
device: cpu-small
|
||||
commands:
|
||||
- pytest -v -s models/test_utils.py models/test_vision.py
|
||||
- pytest -v -s models/test_utils.py models/test_vision.py models/transformers/fusers/
|
||||
|
||||
@@ -17,7 +17,7 @@ steps:
|
||||
- TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'
|
||||
- CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'
|
||||
# Avoid importing model tests that cause CUDA reinitialization error
|
||||
- pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest models/transformers/test_backend.py -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest models/language -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest models/multimodal/generation/test_phi4siglip.py -v -s -m 'distributed(num_gpus=2)'
|
||||
- pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_phi4siglip.py
|
||||
|
||||
@@ -27,6 +27,7 @@ steps:
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- pytest -v -s models/multimodal/generation/test_common.py -m core_model -k "qwen3 or gemma"
|
||||
- pytest -v -s models/multimodal/generation/test_mm_prefix_lm.py -m core_model
|
||||
- pytest -v -s models/multimodal/generation/test_qwen2_5_vl.py -m core_model
|
||||
mirror:
|
||||
amd:
|
||||
@@ -58,7 +59,7 @@ steps:
|
||||
- vllm/
|
||||
- tests/models/multimodal
|
||||
commands:
|
||||
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
|
||||
- pytest -v -s models/multimodal -m core_model --ignore models/multimodal/generation/test_common.py --ignore models/multimodal/generation/test_ultravox.py --ignore models/multimodal/generation/test_qwen2_5_vl.py --ignore models/multimodal/generation/test_qwen2_vl.py --ignore models/multimodal/generation/test_whisper.py --ignore models/multimodal/generation/test_mm_prefix_lm.py --ignore models/multimodal/generation/test_memory_leak.py --ignore models/multimodal/generation/test_vit_cudagraph.py --ignore models/multimodal/processing
|
||||
- pytest -v -s models/multimodal/generation/test_vit_cudagraph.py -m core_model
|
||||
- pytest models/multimodal/generation/test_memory_leak.py -m core_model
|
||||
- cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model # Otherwise, mp_method="spawn" doesn't work
|
||||
|
||||
@@ -15,24 +15,26 @@ steps:
|
||||
- tests/utils.py
|
||||
- tests/benchmarks/test_serve_cli.py
|
||||
- tests/entrypoints/openai/chat_completion/test_chat_completion.py
|
||||
# - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
|
||||
- tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
|
||||
|
||||
# - tests/entrypoints/openai/completion/test_prompt_validation.py
|
||||
- tests/entrypoints/openai/completion/test_shutdown.py
|
||||
# - tests/entrypoints/openai/test_return_token_ids.py
|
||||
# - tests/entrypoints/openai/test_uds.py
|
||||
- tests/entrypoints/openai/test_return_token_ids.py
|
||||
- tests/entrypoints/openai/test_uds.py
|
||||
- tests/v1/sample/test_logprobs_e2e.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex"
|
||||
# - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not invalid"
|
||||
- pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
|
||||
|
||||
# - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
|
||||
- pytest -v -s entrypoints/openai/completion/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
|
||||
# - pytest -v -s entrypoints/openai/test_return_token_ids.py
|
||||
# - pytest -v -s entrypoints/openai/test_uds.py
|
||||
# test_comparison streams differently: Rust emits a separate first (prompt_token_ids) chunk and
|
||||
# finish chunk without logprobs, while the test reads `logprobs.tokens` on every chunk.
|
||||
- pytest -v -s entrypoints/openai/test_return_token_ids.py -k "not test_comparison"
|
||||
- pytest -v -s entrypoints/openai/test_uds.py
|
||||
- pytest -v -s v1/sample/test_logprobs_e2e.py -k "test_prompt_logprobs_e2e_server"
|
||||
|
||||
- label: Rust Frontend Serve/Admin Coverage
|
||||
@@ -45,19 +47,24 @@ steps:
|
||||
- vllm/entrypoints/serve/
|
||||
- vllm/v1/engine/
|
||||
- tests/utils.py
|
||||
# - tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
- tests/entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
- tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py
|
||||
- tests/entrypoints/serve/instrumentator/test_basic.py
|
||||
- tests/entrypoints/serve/instrumentator/test_metrics.py
|
||||
# - tests/entrypoints/serve/dev/test_sleep.py
|
||||
- tests/entrypoints/serve/tokenize/test_tokenization.py
|
||||
commands:
|
||||
- export VLLM_USE_RUST_FRONTEND=1
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
- PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py
|
||||
# server_load can be flaky under the Rust frontend; keep it excluded for now.
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load"
|
||||
# test_generate_logprobs expects Python-style top_logprobs truncation (dedup sampled + cap at max(k, 1)).
|
||||
- pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow"
|
||||
- pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist"
|
||||
# - pytest -v -s entrypoints/serve/dev/test_sleep.py
|
||||
# /tokenizer_info is not implemented in the Rust frontend (the CLI flag is accepted as a no-op).
|
||||
- pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info and not test_tokenize_chat"
|
||||
|
||||
- label: Rust Frontend Core Correctness
|
||||
timeout_in_minutes: 30
|
||||
|
||||
+1
-1
@@ -119,7 +119,7 @@
|
||||
|
||||
# Transformers modeling backend
|
||||
/vllm/model_executor/models/transformers @hmellor
|
||||
/tests/models/test_transformers.py @hmellor
|
||||
/tests/models/transformers @hmellor
|
||||
|
||||
# Docs
|
||||
/docs/mkdocs @hmellor
|
||||
|
||||
@@ -29,6 +29,7 @@ Do not open one-off PRs for tiny edits (single typo, isolated style change, one
|
||||
- PR descriptions for AI-assisted work **must** include:
|
||||
- Why this is not duplicating an existing PR.
|
||||
- Test commands run and results.
|
||||
- Model evaluation results when the change affects output, accuracy, or serving.
|
||||
- Clear statement that AI assistance was used.
|
||||
|
||||
### Fail-closed behavior
|
||||
@@ -66,23 +67,38 @@ VLLM_USE_PRECOMPILED=1 uv pip install -e . --torch-backend=auto
|
||||
uv pip install -e . --torch-backend=auto
|
||||
```
|
||||
|
||||
### Running tests
|
||||
### Tests
|
||||
|
||||
> Requires [Environment setup](#environment-setup) and [Installing dependencies](#installing-dependencies).
|
||||
|
||||
```bash
|
||||
# Install test dependencies.
|
||||
# requirements/test/cuda.txt is pinned to x86_64; on other platforms, use the
|
||||
# unpinned source file instead:
|
||||
uv pip install -r requirements/test/cuda.in # resolves for current platform
|
||||
# Or on x86_64:
|
||||
uv pip install -r requirements/test/cuda.txt
|
||||
# Install test dependencies (use cuda.in on non-x86_64):
|
||||
uv pip install -r requirements/test/cuda.in
|
||||
|
||||
# Run a specific test file (use .venv/bin/python directly;
|
||||
# `source activate` does not persist in non-interactive shells):
|
||||
# Run a specific test file:
|
||||
.venv/bin/python -m pytest tests/path/to/test_file.py -v
|
||||
```
|
||||
|
||||
When adding tests:
|
||||
|
||||
- **Design before you write.** Answer four questions first: what is the module
|
||||
for, what is its I/O contract, what failure am I guarding against, and what is
|
||||
the cheapest level that catches it (unit over integration over e2e)?
|
||||
- **Reuse before create.** Extend existing test files, `conftest.py` fixtures, and
|
||||
helpers; add a new file only when no nearby suite fits.
|
||||
- **Test behavior with intent.** Assert observable outcomes through public APIs;
|
||||
state why in the name or docstring. Skip trivial wiring; flaky tests are worse
|
||||
than no tests.
|
||||
- **Keep it minimal.** One behavior per test and the smallest setup that
|
||||
triggers it; if the test diff dwarfs the code change, cut scope.
|
||||
- **No one-off kernel benchmarks in `tests/`.** Put kernel perf work in
|
||||
`benchmarks/kernels/`; prove correctness in existing pytest suites.
|
||||
- **Run model evals for model-affecting changes.** Search `tests/evals/` or use
|
||||
`vllm bench` and include results in the PR — do not wait for reviewers to ask.
|
||||
|
||||
For model-specific requirements, see
|
||||
[`docs/contributing/model/tests.md`](docs/contributing/model/tests.md).
|
||||
|
||||
### Running linters
|
||||
|
||||
> Requires [Environment setup](#environment-setup).
|
||||
@@ -107,23 +123,18 @@ Use [Google-style docstrings](https://google.github.io/styleguide/pyguide.html#3
|
||||
|
||||
### Coding style guidelines
|
||||
|
||||
Follow these rules for all code changes in this repository:
|
||||
|
||||
- Try to match existing code style.
|
||||
- Code should be self-documenting and self-explanatory.
|
||||
- Keep comments and docstrings minimal and concise.
|
||||
- Match existing code style
|
||||
- Minimize use of comments. Eliminate comments which are redundant, preferring legible and self-documenting code. When used, keep docstrings and comments brief and direct.
|
||||
- Assume the reader is familiar with vLLM.
|
||||
|
||||
### Commit messages
|
||||
|
||||
Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`). For example:
|
||||
Add attribution using commit trailers such as `Co-authored-by:` (other projects use `Assisted-by:` or `Generated-by:`):
|
||||
|
||||
```text
|
||||
Your commit message here
|
||||
|
||||
Co-authored-by: GitHub Copilot
|
||||
Co-authored-by: Claude
|
||||
Co-authored-by: gemini-code-assist
|
||||
Co-authored-by: Agent Name Here
|
||||
Signed-off-by: Your Name <your.email@example.com>
|
||||
```
|
||||
|
||||
|
||||
@@ -400,7 +400,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
|
||||
"csrc/libtorch_stable/topk.cu"
|
||||
"csrc/libtorch_stable/mamba/selective_scan_fwd.cu"
|
||||
"csrc/libtorch_stable/cache_kernels.cu"
|
||||
"csrc/libtorch_stable/cache_kernels.cu"
|
||||
"csrc/libtorch_stable/cache_kernels_fused.cu"
|
||||
"csrc/libtorch_stable/custom_all_reduce.cu"
|
||||
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")
|
||||
|
||||
@@ -12,7 +12,7 @@ from dataclasses import dataclass, field
|
||||
import aiohttp
|
||||
import huggingface_hub.constants
|
||||
from tqdm.asyncio import tqdm
|
||||
from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
|
||||
from transformers import AutoTokenizer, PythonBackend, TokenizersBackend
|
||||
|
||||
# NOTE(simon): do not import vLLM here so the benchmark script
|
||||
# can run without vLLM installed.
|
||||
@@ -609,7 +609,7 @@ def get_tokenizer(
|
||||
tokenizer_mode: str = "auto",
|
||||
trust_remote_code: bool = False,
|
||||
**kwargs,
|
||||
) -> PreTrainedTokenizer | PreTrainedTokenizerFast:
|
||||
) -> PythonBackend | TokenizersBackend:
|
||||
if pretrained_model_name_or_path is not None and not os.path.exists(
|
||||
pretrained_model_name_or_path
|
||||
):
|
||||
|
||||
@@ -132,8 +132,10 @@ def benchmark_function(
|
||||
reset_memory_stats()
|
||||
|
||||
# Benchmark
|
||||
start_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
|
||||
end_events = [torch.Event(enable_timing=True) for _ in range(benchmark_iters)]
|
||||
start_events = [
|
||||
torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)
|
||||
]
|
||||
end_events = [torch.cuda.Event(enable_timing=True) for _ in range(benchmark_iters)]
|
||||
|
||||
for i in range(benchmark_iters):
|
||||
logits_copy = logits.clone()
|
||||
|
||||
@@ -134,8 +134,8 @@ def benchmark_config(
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
# Benchmark
|
||||
start = torch.Event(enable_timing=True)
|
||||
end = torch.Event(enable_timing=True)
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(num_iters):
|
||||
with override_config(config):
|
||||
|
||||
@@ -170,8 +170,8 @@ def benchmark_config(
|
||||
graph.replay()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
start = torch.Event(enable_timing=True)
|
||||
end = torch.Event(enable_timing=True)
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
latencies: list[float] = []
|
||||
for _ in range(num_iters):
|
||||
start.record()
|
||||
|
||||
@@ -15,6 +15,7 @@ endif()
|
||||
#
|
||||
set(ENABLE_X86_ISA $ENV{VLLM_CPU_X86})
|
||||
set(ENABLE_ARM_BF16 $ENV{VLLM_CPU_ARM_BF16})
|
||||
set(ENABLE_RVV_BF16 $ENV{VLLM_CPU_RVV_BF16})
|
||||
|
||||
include_directories("${CMAKE_SOURCE_DIR}/csrc")
|
||||
|
||||
@@ -110,6 +111,13 @@ else()
|
||||
set(ARM_BF16_FOUND ON)
|
||||
message(STATUS "ARM BF16 support enabled via VLLM_CPU_ARM_BF16 environment variable")
|
||||
endif()
|
||||
# Some kernels (e.g. Bianbu on Spacemit X100) do not report zvfbfmin
|
||||
# in /proc/cpuinfo despite hardware support. VLLM_CPU_RVV_BF16=1
|
||||
# overrides the detection result.
|
||||
if (ENABLE_RVV_BF16)
|
||||
set(RVV_BF16_FOUND ON)
|
||||
message(STATUS "RVV BF16 support enabled via VLLM_CPU_RVV_BF16 environment variable")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if (CMAKE_SYSTEM_PROCESSOR MATCHES "x86_64|amd64" OR ENABLE_X86_ISA)
|
||||
@@ -178,7 +186,10 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
|
||||
# Override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256 for RVV.
|
||||
if(NOT DEFINED VLLM_RVV_VLEN)
|
||||
# Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
|
||||
if(EXISTS /proc/cpuinfo)
|
||||
# Skip when cross-compiling — /proc/cpuinfo describes the build host.
|
||||
if(CMAKE_CROSSCOMPILING)
|
||||
message(STATUS "Cross-compiling: skipping VLEN auto-detection from /proc/cpuinfo")
|
||||
elseif(EXISTS /proc/cpuinfo)
|
||||
file(READ /proc/cpuinfo _cpuinfo)
|
||||
set(_best 0)
|
||||
foreach(_n IN ITEMS 128 256 512 1024)
|
||||
@@ -186,6 +197,13 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
|
||||
set(_best ${_n})
|
||||
endif()
|
||||
endforeach()
|
||||
# Only VLEN=128 and VLEN=256 are supported by the RVV kernels.
|
||||
if(_best GREATER 256)
|
||||
message(WARNING
|
||||
"Detected VLEN=${_best} but only 128/256 are supported; "
|
||||
"clamping to 256")
|
||||
set(_best 256)
|
||||
endif()
|
||||
if(_best GREATER 0)
|
||||
set(VLLM_RVV_VLEN ${_best})
|
||||
endif()
|
||||
@@ -195,9 +213,9 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
|
||||
if(NOT DEFINED VLLM_RVV_VLEN AND (RVV_FP16_FOUND OR RVV_BF16_FOUND))
|
||||
message(FATAL_ERROR
|
||||
"RISC-V RVV is available but VLEN could not be auto-detected. "
|
||||
"Please specify VLEN explicitly:\n"
|
||||
" -DVLLM_RVV_VLEN=128 (for VLEN=128 hardware)\n"
|
||||
" -DVLLM_RVV_VLEN=256 (for VLEN=256 hardware, e.g. Spacemit X100)")
|
||||
"Please specify VLEN explicitly via CMAKE_ARGS:\n"
|
||||
" CMAKE_ARGS='-DVLLM_RVV_VLEN=128' (for VLEN=128 hardware)\n"
|
||||
" CMAKE_ARGS='-DVLLM_RVV_VLEN=256' (for VLEN=256 hardware, e.g. Spacemit X100)")
|
||||
endif()
|
||||
endif()
|
||||
if(VLLM_RVV_VLEN AND VLLM_RVV_VLEN GREATER 0)
|
||||
@@ -209,7 +227,7 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
|
||||
message(STATUS "BF16 extension detected")
|
||||
set(MARCH_FLAGS -march=rv64gcv_zvfh_zfbfmin_zvfbfmin_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
|
||||
elseif(RVV_FP16_FOUND)
|
||||
message(WARNING "BF16 functionality is not available")
|
||||
message(WARNING "BF16 functionality is not available.")
|
||||
set(MARCH_FLAGS -march=rv64gcv_zvfh_zvl${VLLM_RVV_VLEN}b -mrvv-vector-bits=zvl -mabi=lp64d)
|
||||
else()
|
||||
message(STATUS "compile riscv with scalar (no FP16/BF16)")
|
||||
|
||||
@@ -13,7 +13,8 @@ static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
|
||||
|
||||
bool cpu_attn_has_isa(const std::string& isa) {
|
||||
if (isa == "rvv") {
|
||||
#if defined(__riscv) && defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128
|
||||
#if defined(__riscv) && defined(__riscv_v_min_vlen) && \
|
||||
(__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
|
||||
@@ -323,8 +323,6 @@ class AttentionImpl<ISA::VSX, scalar_t, head_dim> {
|
||||
const int64_t num_blocks_stride, const int64_t cache_head_num_stride,
|
||||
const int64_t block_size, const int64_t block_size_stride,
|
||||
const float k_inv = 0.0f, const float v_inv = 0.0f) {
|
||||
// k_inv and v_inv are unused on VSX: FP8 KV cache is not supported on
|
||||
// PowerPC. The parameters are present to match the common interface.
|
||||
#pragma omp parallel for collapse(2)
|
||||
for (int64_t token_idx = 0; token_idx < token_num; ++token_idx) {
|
||||
for (int64_t head_idx = 0; head_idx < head_num; ++head_idx) {
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
#if defined(__x86_64__)
|
||||
// x86 implementation
|
||||
#include "cpu_types_x86.hpp"
|
||||
#elif defined(__POWER9_VECTOR__)
|
||||
#elif defined(__powerpc__)
|
||||
// ppc implementation
|
||||
#include "cpu_types_vsx.hpp"
|
||||
#elif defined(__s390x__)
|
||||
@@ -41,4 +41,4 @@ inline int get_max_threads() {
|
||||
}
|
||||
} // namespace cpu_utils
|
||||
|
||||
#endif
|
||||
#endif
|
||||
|
||||
@@ -214,11 +214,18 @@ struct BF16Vec32 : public Vec<BF16Vec32> {
|
||||
|
||||
explicit BF16Vec32(const BF16Vec8& v) {
|
||||
fixed_u16x8_t u16_val = bf16_to_u16(v.reg);
|
||||
fixed_u16x32_t u16_combined =
|
||||
RVVI4(__riscv_vcreate_v_u16, LMUL_128, _u16, LMUL_512)(
|
||||
u16_val, u16_val, u16_val, u16_val);
|
||||
reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16,
|
||||
LMUL_512)(u16_combined);
|
||||
// Widen LMUL_128 → LMUL_256 so vslideup operands share a type.
|
||||
// At VLEN=256 this is mf2→m1 (both integer); at VLEN=128 it is m1→m2.
|
||||
fixed_u16x16_t ext =
|
||||
RVVI4(__riscv_vlmul_ext_v_u16, LMUL_128, _u16, LMUL_256)(u16_val);
|
||||
// Build 16-element half: place the 8 elements at offsets 0 and 8.
|
||||
fixed_u16x16_t half = RVVI(__riscv_vmv_v_x_u16, LMUL_256)(0, 16);
|
||||
half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 0, 8);
|
||||
half = RVVI(__riscv_vslideup_vx_u16, LMUL_256)(half, ext, 8, 16);
|
||||
// Double to LMUL_512 (m1→m2 at VLEN=256, m2→m4 at VLEN=128).
|
||||
fixed_u16x32_t dst =
|
||||
RVVI4(__riscv_vcreate_v_u16, LMUL_256, _u16, LMUL_512)(half, half);
|
||||
reg = RVVI4(__riscv_vreinterpret_v_u16, LMUL_512, _bf16, LMUL_512)(dst);
|
||||
};
|
||||
|
||||
void save(void* ptr) const {
|
||||
@@ -623,17 +630,29 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
|
||||
data.reg, data.reg)) {};
|
||||
explicit FP32Vec16(const FP32Vec16& data) : reg(data.reg) {};
|
||||
explicit FP32Vec16(int64_t value, const FP32Vec16& lut) {
|
||||
const uint64_t q_values = static_cast<uint64_t>(value);
|
||||
auto packed = RVVI(__riscv_vmv_v_x_u64, LMUL_1024)(q_values, VEC_ELEM_NUM);
|
||||
auto lane_ids = RVVI(__riscv_vid_v_u64, LMUL_1024)(VEC_ELEM_NUM);
|
||||
auto shifts =
|
||||
RVVI(__riscv_vsll_vx_u64, LMUL_1024)(lane_ids, 2, VEC_ELEM_NUM);
|
||||
auto shifted =
|
||||
RVVI(__riscv_vsrl_vv_u64, LMUL_1024)(packed, shifts, VEC_ELEM_NUM);
|
||||
auto idx64 =
|
||||
RVVI(__riscv_vand_vx_u64, LMUL_1024)(shifted, 0xF, VEC_ELEM_NUM);
|
||||
auto idx32 = RVVI(__riscv_vnsrl_wx_u32, LMUL_512)(idx64, 0, VEC_ELEM_NUM);
|
||||
reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx32, VEC_ELEM_NUM);
|
||||
// Split into two 32-bit halves to avoid u64 @ LMUL_1024 (m8 on
|
||||
// VLEN=128 / m4 on VLEN=256), which causes heavy register spilling.
|
||||
constexpr int HALF = VEC_ELEM_NUM / 2;
|
||||
const auto q = static_cast<uint64_t>(value);
|
||||
const uint32_t lo = static_cast<uint32_t>(q);
|
||||
const uint32_t hi = static_cast<uint32_t>(q >> 32);
|
||||
|
||||
auto lane_ids = RVVI(__riscv_vid_v_u32, LMUL_256)(HALF);
|
||||
auto shifts = RVVI(__riscv_vsll_vx_u32, LMUL_256)(lane_ids, 2, HALF);
|
||||
|
||||
auto packed_lo = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(lo, HALF);
|
||||
auto idx_lo = RVVI(__riscv_vand_vx_u32, LMUL_256)(
|
||||
RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_lo, shifts, HALF), 0xF,
|
||||
HALF);
|
||||
|
||||
auto packed_hi = RVVI(__riscv_vmv_v_x_u32, LMUL_256)(hi, HALF);
|
||||
auto idx_hi = RVVI(__riscv_vand_vx_u32, LMUL_256)(
|
||||
RVVI(__riscv_vsrl_vv_u32, LMUL_256)(packed_hi, shifts, HALF), 0xF,
|
||||
HALF);
|
||||
|
||||
auto idx =
|
||||
RVVI4(__riscv_vcreate_v_u32, LMUL_256, _u32, LMUL_512)(idx_lo, idx_hi);
|
||||
reg = RVVI(__riscv_vrgather_vv_f32, LMUL_512)(lut.reg, idx, VEC_ELEM_NUM);
|
||||
}
|
||||
explicit FP32Vec16(const FP16Vec16& v);
|
||||
|
||||
|
||||
+120
-40
@@ -344,53 +344,133 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
FP32Vec8 exp() const {
|
||||
// TODO: Vectorize this
|
||||
AliasReg ar;
|
||||
ar.reg = reg;
|
||||
f32x4x4_t ret;
|
||||
ret.val[0][0] = std::exp(ar.values[0]);
|
||||
ret.val[0][1] = std::exp(ar.values[1]);
|
||||
ret.val[0][2] = std::exp(ar.values[2]);
|
||||
ret.val[0][3] = std::exp(ar.values[3]);
|
||||
ret.val[1][0] = std::exp(ar.values[4]);
|
||||
ret.val[1][1] = std::exp(ar.values[5]);
|
||||
ret.val[1][2] = std::exp(ar.values[6]);
|
||||
ret.val[1][3] = std::exp(ar.values[7]);
|
||||
return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
|
||||
f32x4x2_t out;
|
||||
const __vector float log2e = vec_splats(1.44269504088896341f);
|
||||
const __vector float one = vec_splats(1.0f);
|
||||
const __vector float min_x = vec_splats(-87.3f);
|
||||
const __vector float max_x = vec_splats(88.7f);
|
||||
|
||||
// 5th-degree minimax polynomial for 2^r (r in [0,1))
|
||||
const __vector float c1 = vec_splats(0.6931471805599453f);
|
||||
const __vector float c2 = vec_splats(0.240226506959101f);
|
||||
const __vector float c3 = vec_splats(0.05550410866482158f);
|
||||
const __vector float c4 = vec_splats(0.009618129107628477f);
|
||||
const __vector float c5 = vec_splats(0.0013333558146428443f);
|
||||
|
||||
for (int i = 0; i < 2; i++) {
|
||||
__vector float x = reg.val[i];
|
||||
x = vec_max(x, min_x);
|
||||
x = vec_min(x, max_x);
|
||||
|
||||
__vector float y = vec_mul(x, log2e);
|
||||
|
||||
__vector float kf = vec_floor(y);
|
||||
__vector float r = vec_sub(y, kf);
|
||||
|
||||
// Convert float to signed integer. Use vec_cts for PowerPC AltiVec
|
||||
// compatibility.
|
||||
__vector signed int k = vec_cts(kf, 0);
|
||||
const __vector signed int min_k = vec_splats((signed int)-126);
|
||||
const __vector signed int max_k = vec_splats((signed int)127);
|
||||
k = vec_min(vec_max(k, min_k), max_k);
|
||||
|
||||
// Build 2^k from exponent bits
|
||||
__vector signed int exp_int = vec_add(k, vec_splats((signed int)127));
|
||||
__vector unsigned int bits = (__vector unsigned int)exp_int;
|
||||
bits = vec_sl(bits, vec_splats((unsigned int)23));
|
||||
__vector float pow2k = (__vector float)bits;
|
||||
|
||||
// Improved minimax polynomial
|
||||
__vector float poly = vec_madd(c5, r, c4);
|
||||
poly = vec_madd(poly, r, c3);
|
||||
poly = vec_madd(poly, r, c2);
|
||||
poly = vec_madd(poly, r, c1);
|
||||
poly = vec_madd(poly, r, one);
|
||||
|
||||
out.val[i] = vec_mul(pow2k, poly);
|
||||
}
|
||||
return FP32Vec8(out);
|
||||
}
|
||||
|
||||
FP32Vec8 tanh() const {
|
||||
// TODO: Vectorize this
|
||||
AliasReg ar;
|
||||
ar.reg = reg;
|
||||
f32x4x4_t ret;
|
||||
ret.val[0][0] = std::tanh(ar.values[0]);
|
||||
ret.val[0][1] = std::tanh(ar.values[1]);
|
||||
ret.val[0][2] = std::tanh(ar.values[2]);
|
||||
ret.val[0][3] = std::tanh(ar.values[3]);
|
||||
ret.val[1][0] = std::tanh(ar.values[4]);
|
||||
ret.val[1][1] = std::tanh(ar.values[5]);
|
||||
ret.val[1][2] = std::tanh(ar.values[6]);
|
||||
ret.val[1][3] = std::tanh(ar.values[7]);
|
||||
return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
|
||||
const __vector float one = vec_splats(1.0f);
|
||||
const __vector float two = vec_splats(2.0f);
|
||||
const __vector float zero = vec_splats(0.0f);
|
||||
const __vector float sat = vec_splats(9.0f);
|
||||
|
||||
f32x4x2_t out;
|
||||
|
||||
for (int i = 0; i < 2; i++) {
|
||||
__vector float x = reg.val[i];
|
||||
__vector float ax = vec_abs(x);
|
||||
|
||||
__vector bool int mask = vec_cmpge(x, zero);
|
||||
__vector float sign = vec_sel(vec_splats(-1.0f), one, mask);
|
||||
|
||||
__vector bool int saturated = vec_cmpge(ax, sat);
|
||||
|
||||
__vector float two_x = vec_mul(x, two);
|
||||
f32x4x2_t tmp;
|
||||
tmp.val[0] = two_x;
|
||||
tmp.val[1] = two_x;
|
||||
FP32Vec8 temp_vec(tmp);
|
||||
vector float e = temp_vec.exp().reg.val[0];
|
||||
|
||||
vector float num = vec_sub(e, one);
|
||||
vector float den = vec_add(e, one);
|
||||
vector float t = vec_div(num, den);
|
||||
|
||||
out.val[i] = vec_sel(t, sign, saturated);
|
||||
}
|
||||
return FP32Vec8(out);
|
||||
}
|
||||
|
||||
FP32Vec8 er() const {
|
||||
// TODO: Vectorize this
|
||||
AliasReg ar;
|
||||
ar.reg = reg;
|
||||
f32x4x4_t ret;
|
||||
ret.val[0][0] = std::erf(ar.values[0]);
|
||||
ret.val[0][1] = std::erf(ar.values[1]);
|
||||
ret.val[0][2] = std::erf(ar.values[2]);
|
||||
ret.val[0][3] = std::erf(ar.values[3]);
|
||||
ret.val[1][0] = std::erf(ar.values[4]);
|
||||
ret.val[1][1] = std::erf(ar.values[5]);
|
||||
ret.val[1][2] = std::erf(ar.values[6]);
|
||||
ret.val[1][3] = std::erf(ar.values[7]);
|
||||
return FP32Vec8(f32x4x2_t({ret.val[0], ret.val[1]}));
|
||||
const vector float a1 = vec_splats(0.254829592f);
|
||||
const vector float a2 = vec_splats(-0.284496736f);
|
||||
const vector float a3 = vec_splats(1.421413741f);
|
||||
const vector float a4 = vec_splats(-1.453152027f);
|
||||
const vector float a5 = vec_splats(1.061405429f);
|
||||
const vector float p = vec_splats(0.3275911f);
|
||||
const vector float one = vec_splats(1.0f);
|
||||
const vector float zero = vec_splats(0.0f);
|
||||
const vector float sat = vec_splats(6.0f);
|
||||
|
||||
f32x4x2_t ret;
|
||||
|
||||
for (int i = 0; i < 2; i++) {
|
||||
vector float x = reg.val[i];
|
||||
vector float ax = vec_abs(x);
|
||||
|
||||
vector bool int mask = vec_cmpge(x, zero);
|
||||
vector float sign = vec_sel(vec_splats(-1.0f), one, mask);
|
||||
|
||||
vector bool int saturated = vec_cmpge(ax, sat);
|
||||
|
||||
vector float t = vec_div(one, vec_madd(p, ax, one));
|
||||
|
||||
vector float poly = a5;
|
||||
poly = vec_madd(poly, t, a4);
|
||||
poly = vec_madd(poly, t, a3);
|
||||
poly = vec_madd(poly, t, a2);
|
||||
poly = vec_madd(poly, t, a1);
|
||||
poly = vec_mul(poly, t);
|
||||
|
||||
vector float x_squared = vec_mul(x, x);
|
||||
vector float neg_x_squared = vec_mul(vec_splats(-1.0f), x_squared);
|
||||
f32x4x2_t tmp;
|
||||
tmp.val[0] = neg_x_squared;
|
||||
tmp.val[1] = neg_x_squared;
|
||||
FP32Vec8 exp_input(tmp);
|
||||
vector float exp_term = exp_input.exp().reg.val[0];
|
||||
|
||||
vector float y = vec_nmsub(poly, exp_term, one);
|
||||
vector float erf_val = vec_mul(sign, y);
|
||||
|
||||
ret.val[i] = vec_sel(erf_val, sign, saturated);
|
||||
}
|
||||
return FP32Vec8(ret);
|
||||
}
|
||||
|
||||
FP32Vec8 operator*(const FP32Vec8& b) const {
|
||||
|
||||
@@ -278,7 +278,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
|
||||
ops.def(
|
||||
"dynamic_4bit_int_moe("
|
||||
"Tensor x, Tensor topk_ids, Tensor topk_weights,"
|
||||
"Tensor w13_packed, Tensor w2_packed, int H, int I, int I2,"
|
||||
"Tensor w13_packed, Tensor w2_packed,"
|
||||
"int hidden_size, int intermediate_size,"
|
||||
"int group_size, bool apply_router_weight_on_input, int activation_kind"
|
||||
") -> Tensor");
|
||||
|
||||
|
||||
+2
-2
@@ -76,14 +76,14 @@ inline int64_t get_available_l2_size() {
|
||||
if (l2_cache_size == 0) {
|
||||
l2_cache_size = 256 * 1024;
|
||||
}
|
||||
return static_cast<int64_t>(l2_cache_size) >> 1; // use 50% of L2 cache
|
||||
return static_cast<int64_t>(l2_cache_size) >> 1;
|
||||
}();
|
||||
return size;
|
||||
#else
|
||||
static int64_t size = []() {
|
||||
auto caps = at::cpu::get_cpu_capabilities();
|
||||
const uint32_t l2_cache_size = caps.at("l2_cache_size").toInt();
|
||||
return l2_cache_size >> 1; // use 50% of L2 cache
|
||||
return l2_cache_size >> 1;
|
||||
}();
|
||||
return size;
|
||||
#endif
|
||||
|
||||
@@ -15,7 +15,8 @@ void topk_sigmoid(torch::stable::Tensor& topk_weights,
|
||||
torch::stable::Tensor& topk_indices,
|
||||
torch::stable::Tensor& token_expert_indices,
|
||||
torch::stable::Tensor& gating_output, bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias);
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor);
|
||||
|
||||
void topk_softplus_sqrt(
|
||||
torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
|
||||
|
||||
@@ -173,7 +173,8 @@ __launch_bounds__(TPB) __global__ void moeTopK(
|
||||
const int start_expert,
|
||||
const int end_expert,
|
||||
const bool renormalize,
|
||||
const float* bias)
|
||||
const float* bias,
|
||||
const double routed_scaling_factor)
|
||||
{
|
||||
|
||||
using cub_kvp = cub::KeyValuePair<int, float>;
|
||||
@@ -241,14 +242,16 @@ __launch_bounds__(TPB) __global__ void moeTopK(
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Renormalize the k weights for this row to sum to 1, if requested.
|
||||
if (renormalize) {
|
||||
if (threadIdx.x == 0) {
|
||||
// Apply renormalization and routed scaling factor to final weights.
|
||||
if (threadIdx.x == 0) {
|
||||
float scale = static_cast<float>(routed_scaling_factor);
|
||||
if (renormalize) {
|
||||
const float denom = selected_sum > 0.f ? selected_sum : 1.f;
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int idx = k * block_row + k_idx;
|
||||
output[idx] = output[idx] / denom;
|
||||
}
|
||||
scale /= denom;
|
||||
}
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int idx = k * block_row + k_idx;
|
||||
output[idx] = output[idx] * scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -274,7 +277,7 @@ template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WA
|
||||
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
|
||||
int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
|
||||
const float* bias)
|
||||
const float* bias, const double routed_scaling_factor)
|
||||
{
|
||||
static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
|
||||
std::is_same_v<InputType, __half>,
|
||||
@@ -570,17 +573,17 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
|
||||
}
|
||||
}
|
||||
|
||||
// Renormalize the k weights for this row to sum to 1, if requested.
|
||||
if (renormalize) {
|
||||
if (thread_group_idx == 0)
|
||||
{
|
||||
const float denom = selected_sum > 0.f ? selected_sum : 1.f;
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx)
|
||||
{
|
||||
const int idx = k * thread_row + k_idx;
|
||||
output[idx] = output[idx] / denom;
|
||||
}
|
||||
}
|
||||
// Apply renormalization and routed scaling factor to final weights.
|
||||
if (thread_group_idx == 0) {
|
||||
float scale = static_cast<float>(routed_scaling_factor);
|
||||
if (renormalize) {
|
||||
const float denom = selected_sum > 0.f ? selected_sum : 1.f;
|
||||
scale /= denom;
|
||||
}
|
||||
for (int k_idx = 0; k_idx < k; ++k_idx) {
|
||||
const int idx = k * thread_row + k_idx;
|
||||
output[idx] = output[idx] * scale;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -602,7 +605,7 @@ struct TopkConstants
|
||||
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
|
||||
void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices,
|
||||
int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
|
||||
const float* bias, cudaStream_t stream)
|
||||
const float* bias, const double routed_scaling_factor, cudaStream_t stream)
|
||||
{
|
||||
static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
|
||||
using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
|
||||
@@ -613,7 +616,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
|
||||
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
|
||||
topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias);
|
||||
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);
|
||||
}
|
||||
|
||||
#ifndef USE_ROCM
|
||||
@@ -624,7 +627,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, stream);
|
||||
bias, routed_scaling_factor, stream);
|
||||
#else
|
||||
#define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
|
||||
if (WARP_SIZE == 64) { \
|
||||
@@ -632,13 +635,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, stream); \
|
||||
bias, routed_scaling_factor, stream); \
|
||||
} else if (WARP_SIZE == 32) { \
|
||||
topkGatingLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES, \
|
||||
IndType, InputType, SF>( \
|
||||
gating_output, nullptr, topk_weights, topk_indices, \
|
||||
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
|
||||
bias, stream); \
|
||||
bias, routed_scaling_factor, stream); \
|
||||
} else { \
|
||||
assert(false && \
|
||||
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
|
||||
@@ -657,6 +660,7 @@ void topkGatingKernelLauncher(
|
||||
const int topk,
|
||||
const bool renormalize,
|
||||
const float* bias,
|
||||
const double routed_scaling_factor,
|
||||
cudaStream_t stream) {
|
||||
static constexpr int WARPS_PER_TB = 4;
|
||||
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
|
||||
@@ -732,7 +736,7 @@ void topkGatingKernelLauncher(
|
||||
}
|
||||
moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
|
||||
workspace, nullptr, topk_weights, topk_indices, token_expert_indices,
|
||||
num_experts, topk, 0, num_experts, renormalize, bias);
|
||||
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -750,6 +754,7 @@ void dispatch_topk_launch(
|
||||
torch::stable::Tensor& softmax_workspace,
|
||||
int num_tokens, int num_experts, int topk, bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
const float* bias_ptr = nullptr;
|
||||
@@ -772,7 +777,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, stream);
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
} else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) {
|
||||
vllm::moe::topkGatingKernelLauncher<uint32_t, ComputeType, SF>(
|
||||
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
|
||||
@@ -781,7 +786,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, stream);
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
} else {
|
||||
STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long);
|
||||
vllm::moe::topkGatingKernelLauncher<int64_t, ComputeType, SF>(
|
||||
@@ -791,7 +796,7 @@ void dispatch_topk_launch(
|
||||
token_expert_indices.mutable_data_ptr<int>(),
|
||||
softmax_workspace.mutable_data_ptr<float>(),
|
||||
num_tokens, num_experts, topk, renormalize,
|
||||
bias_ptr, stream);
|
||||
bias_ptr, routed_scaling_factor, stream);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -820,15 +825,15 @@ void topk_softmax(
|
||||
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
|
||||
dispatch_topk_launch<float, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, stream);
|
||||
bias, 1.0, stream);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
|
||||
dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, stream);
|
||||
bias, 1.0, stream);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, stream);
|
||||
bias, 1.0, stream);
|
||||
} else {
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
|
||||
}
|
||||
@@ -840,7 +845,8 @@ void topk_sigmoid(
|
||||
torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
|
||||
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
|
||||
bool renormalize,
|
||||
std::optional<torch::stable::Tensor> bias)
|
||||
std::optional<torch::stable::Tensor> bias,
|
||||
double routed_scaling_factor)
|
||||
{
|
||||
const int num_experts = gating_output.size(-1);
|
||||
const auto num_tokens = gating_output.numel() / num_experts;
|
||||
@@ -859,15 +865,15 @@ void topk_sigmoid(
|
||||
if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) {
|
||||
dispatch_topk_launch<float, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, stream);
|
||||
bias, routed_scaling_factor, stream);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) {
|
||||
dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, stream);
|
||||
bias, routed_scaling_factor, stream);
|
||||
} else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
|
||||
dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices,
|
||||
token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize,
|
||||
bias, stream);
|
||||
bias, routed_scaling_factor, stream);
|
||||
} else {
|
||||
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
|
||||
}
|
||||
|
||||
@@ -13,8 +13,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
|
||||
// Apply topk sigmoid to the gating outputs.
|
||||
m.def(
|
||||
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
|
||||
"bias) -> ()");
|
||||
"token_expert_indices, Tensor gating_output, bool renormalize, "
|
||||
"Tensor? bias, float routed_scaling_factor) -> ()");
|
||||
|
||||
m.def(
|
||||
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
|
||||
|
||||
@@ -29,25 +29,37 @@ enum ActivationKind : int64_t {
|
||||
|
||||
torch::Tensor dynamic_4bit_int_moe_cpu(
|
||||
torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
|
||||
torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
|
||||
int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
|
||||
int64_t activation_kind) {
|
||||
torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size,
|
||||
int64_t intermediate_size, int64_t group_size,
|
||||
bool apply_router_weight_on_input, int64_t activation_kind) {
|
||||
TORCH_CHECK(x.dim() == 2, "x must be 2D");
|
||||
TORCH_CHECK(topk_ids.dim() == 2 && topk_weights.dim() == 2,
|
||||
"topk tensors must be [T, K]");
|
||||
TORCH_CHECK(
|
||||
w13_packed.size(0) == w2_packed.size(0),
|
||||
"w13_packed and w2_packed must have same number of experts in dim 0");
|
||||
TORCH_CHECK(I2 == 2 * I, "I2 must equal 2*I");
|
||||
|
||||
const int64_t T = x.size(0);
|
||||
const int64_t K = topk_ids.size(1);
|
||||
const int64_t E = w13_packed.size(0);
|
||||
const int64_t N = T * K;
|
||||
const int64_t w13_out_features = 2 * intermediate_size;
|
||||
|
||||
auto x_c = x.contiguous();
|
||||
// _dyn_quant_matmul_4bit kernel natively supports these pre-quant activation
|
||||
// dtypes:
|
||||
// - fp32: with channelwise and groupwise
|
||||
// - bf16: with channelwise -> upcast to fp32 for groupwise
|
||||
// - fp16: not supported -> upcast to fp32 for groupwise & channelwise
|
||||
const auto output_dtype = x_c.scalar_type();
|
||||
const bool should_cast_input =
|
||||
((group_size != -1) && output_dtype == at::kBFloat16) ||
|
||||
output_dtype == at::kHalf;
|
||||
if (should_cast_input) {
|
||||
x_c = x_c.to(at::kFloat);
|
||||
}
|
||||
auto ids_c = topk_ids.contiguous();
|
||||
auto gates_c = topk_weights.to(at::kFloat).contiguous();
|
||||
auto gates_c = topk_weights.to(x_c.scalar_type()).contiguous();
|
||||
|
||||
// bucketing tokens -> experts
|
||||
c10::SmallVector<int64_t, 64> counts(
|
||||
@@ -63,35 +75,42 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
|
||||
c10::SmallVector<int64_t, 65> offsets(E + 1, 0); // ( E +1 )
|
||||
for (int64_t e = 0; e < E; ++e) offsets[e + 1] = offsets[e] + counts[e];
|
||||
|
||||
// expert_tokens = [tokens indices for expert 0, ...]
|
||||
// expert_gates = [router weights for tokens assigned to expert 0, ...]
|
||||
auto expert_tokens = at::empty({offsets[E]}, ids_c.options());
|
||||
auto expert_gates = at::empty({offsets[E]}, gates_c.options());
|
||||
{
|
||||
c10::SmallVector<int64_t, 64> cursor(E, 0);
|
||||
const auto* ids_ptr = ids_c.data_ptr<int64_t>();
|
||||
const auto* gts_ptr = gates_c.data_ptr<float>();
|
||||
auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
|
||||
auto* gate_ptr = expert_gates.data_ptr<float>();
|
||||
AT_DISPATCH_FLOATING_TYPES_AND2(
|
||||
at::ScalarType::BFloat16, at::ScalarType::Half, gates_c.scalar_type(),
|
||||
"bucket_expert_tokens_and_gates", [&] {
|
||||
const auto* ids_ptr = ids_c.data_ptr<int64_t>();
|
||||
const auto* gts_ptr = gates_c.data_ptr<scalar_t>();
|
||||
auto* tok_ptr = expert_tokens.data_ptr<int64_t>();
|
||||
auto* gate_ptr = expert_gates.data_ptr<scalar_t>();
|
||||
|
||||
for (int64_t t = 0; t < T; ++t) {
|
||||
const int64_t base = t * K;
|
||||
for (int64_t k = 0; k < K; ++k) {
|
||||
const int64_t idx = base + k;
|
||||
const int64_t e = ids_ptr[idx];
|
||||
const int64_t p = offsets[e] + (cursor[e]++);
|
||||
tok_ptr[p] = t;
|
||||
gate_ptr[p] = gts_ptr[idx];
|
||||
}
|
||||
}
|
||||
for (int64_t t = 0; t < T; ++t) {
|
||||
const int64_t base = t * K;
|
||||
for (int64_t k = 0; k < K; ++k) {
|
||||
const int64_t idx = base + k;
|
||||
const int64_t e = ids_ptr[idx];
|
||||
const int64_t p = offsets[e] + (cursor[e]++);
|
||||
tok_ptr[p] = t;
|
||||
gate_ptr[p] = gts_ptr[idx];
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
const int64_t g_eff_13 = (group_size != -1) ? group_size : H;
|
||||
const int64_t g_eff_2 = (group_size != -1) ? group_size : I;
|
||||
const int64_t g_eff_13 = (group_size != -1) ? group_size : hidden_size;
|
||||
const int64_t g_eff_2 = (group_size != -1) ? group_size : intermediate_size;
|
||||
|
||||
// X_all [num_tokens * K, hidden_size]
|
||||
auto X_all = x_c.index_select(/*dim=*/0, expert_tokens);
|
||||
if (apply_router_weight_on_input) {
|
||||
X_all = X_all.mul(expert_gates.unsqueeze(1));
|
||||
}
|
||||
auto Y_all = at::empty({offsets[E], H}, x_c.options());
|
||||
auto Y_all = at::empty({offsets[E], hidden_size}, x_c.options());
|
||||
|
||||
at::parallel_for(0, offsets[E], 0, [&](int64_t idx_begin, int64_t idx_end) {
|
||||
c10::InferenceMode guard;
|
||||
@@ -109,11 +128,13 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
|
||||
auto w2_e = w2_packed.select(/*dim=*/0, e);
|
||||
|
||||
// W13
|
||||
auto y13 =
|
||||
mm(x_e, w13_e, g_eff_13, /*in_features=*/H, /*out_features=*/I2);
|
||||
auto y13 = mm(x_e, w13_e, g_eff_13, /*in_features=*/hidden_size,
|
||||
/*out_features=*/w13_out_features);
|
||||
|
||||
auto g_part = y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/I);
|
||||
auto u_part = y13.narrow(/*dim=*/1, /*start=*/I, /*length=*/I);
|
||||
auto g_part =
|
||||
y13.narrow(/*dim=*/1, /*start=*/0, /*length=*/intermediate_size);
|
||||
auto u_part = y13.narrow(/*dim=*/1, /*start=*/intermediate_size,
|
||||
/*length=*/intermediate_size);
|
||||
|
||||
torch::Tensor act;
|
||||
if (activation_kind == ActivationKind::SwiGLUOAI) { // SwiGLUOAI
|
||||
@@ -128,7 +149,8 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
|
||||
}
|
||||
|
||||
// W2
|
||||
auto y = mm(act, w2_e, g_eff_2, /*in_features=*/I, /*out_features=*/H);
|
||||
auto y = mm(act, w2_e, g_eff_2, /*in_features=*/intermediate_size,
|
||||
/*out_features=*/hidden_size);
|
||||
|
||||
// Store per-expert result
|
||||
Y_all.narrow(/*dim=*/0, /*start=*/start, /*length=*/te).copy_(y);
|
||||
@@ -138,8 +160,11 @@ torch::Tensor dynamic_4bit_int_moe_cpu(
|
||||
if (!apply_router_weight_on_input) {
|
||||
Y_all = Y_all.mul(expert_gates.unsqueeze(1));
|
||||
}
|
||||
if (Y_all.scalar_type() != output_dtype) {
|
||||
Y_all = Y_all.to(output_dtype);
|
||||
}
|
||||
|
||||
auto out = at::zeros({T, H}, x.options());
|
||||
auto out = at::zeros({T, hidden_size}, x.options());
|
||||
out =
|
||||
at::index_add(out, /*dim=*/0, /*index=*/expert_tokens, /*source=*/Y_all);
|
||||
|
||||
|
||||
+3
-3
@@ -53,9 +53,9 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
|
||||
|
||||
torch::Tensor dynamic_4bit_int_moe_cpu(
|
||||
torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights,
|
||||
torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I,
|
||||
int64_t I2, int64_t group_size, bool apply_router_weight_on_input,
|
||||
int64_t activation_kind);
|
||||
torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t hidden_size,
|
||||
int64_t intermediate_size, int64_t group_size,
|
||||
bool apply_router_weight_on_input, int64_t activation_kind);
|
||||
|
||||
using fptr_t = int64_t;
|
||||
#ifdef USE_ROCM
|
||||
|
||||
@@ -25,7 +25,6 @@ FROM ubuntu:22.04 AS base-common
|
||||
WORKDIR /workspace
|
||||
|
||||
ARG PYTHON_VERSION=3.12
|
||||
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
|
||||
|
||||
ARG max_jobs=32
|
||||
ENV MAX_JOBS=${max_jobs}
|
||||
@@ -53,8 +52,6 @@ ENV PATH="$VIRTUAL_ENV/bin:$PATH"
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# Install Python dependencies
|
||||
ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE="copy"
|
||||
|
||||
@@ -64,7 +61,7 @@ COPY requirements/cpu.txt requirements/cpu.txt
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --upgrade pip && \
|
||||
uv pip install -r requirements/cpu.txt
|
||||
uv pip install -r requirements/cpu.txt --torch-backend cpu
|
||||
|
||||
ARG TARGETARCH
|
||||
ENV TARGETARCH=${TARGETARCH}
|
||||
@@ -149,7 +146,7 @@ RUN if [ "$TARGETARCH" = "arm64" ] && [ "$VLLM_CPU_X86" != "0" ]; then \
|
||||
COPY requirements/build/cpu.txt requirements/build/cpu.txt
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -r requirements/build/cpu.txt
|
||||
uv pip install -r requirements/build/cpu.txt --torch-backend cpu
|
||||
|
||||
COPY . .
|
||||
|
||||
@@ -205,7 +202,7 @@ RUN case "$(uname -m)" in \
|
||||
esac
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -r requirements/test/cpu.txt
|
||||
uv pip install -r requirements/test/cpu.txt --torch-backend cpu
|
||||
|
||||
######################### DEV IMAGE #########################
|
||||
FROM vllm-build AS vllm-dev
|
||||
@@ -231,7 +228,7 @@ COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirement
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -r requirements/lint.txt && \
|
||||
uv pip install -r requirements/test/cpu.txt && \
|
||||
uv pip install -r requirements/test/cpu.txt --torch-backend cpu && \
|
||||
pre-commit install --hook-type pre-commit --hook-type commit-msg
|
||||
|
||||
ENTRYPOINT ["bash"]
|
||||
|
||||
@@ -252,6 +252,8 @@ COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake.hcl /docker/doc
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust /rust
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
|
||||
COPY --from=build_vllm ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
|
||||
|
||||
# RIXL/UCX build stages
|
||||
@@ -543,6 +545,8 @@ COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake.h
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/docker/docker-bake-rocm.hcl /docker/docker-bake-rocm.hcl
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/pyproject.toml /pyproject.toml
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust /rust
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/rust-toolchain.toml /rust-toolchain.toml
|
||||
COPY --from=build_vllm_wheel_release ${COMMON_WORKDIR}/vllm/vllm/v1 /vllm_v1
|
||||
|
||||
# -----------------------
|
||||
@@ -576,6 +580,7 @@ RUN apt-get update -q -y && apt-get install -q -y --no-install-recommends \
|
||||
libibverbs1 \
|
||||
ibverbs-providers \
|
||||
ibverbs-utils \
|
||||
unzip \
|
||||
pkg-config ffmpeg libavcodec-dev libavformat-dev libavutil-dev \
|
||||
libswscale-dev libavdevice-dev libavfilter-dev libswresample-dev \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
@@ -9,7 +9,7 @@ ARG PYTORCH_AUDIO_BRANCH="v2.9.0"
|
||||
ARG PYTORCH_AUDIO_REPO="https://github.com/pytorch/audio.git"
|
||||
ARG FA_BRANCH="0e60e394"
|
||||
ARG FA_REPO="https://github.com/Dao-AILab/flash-attention.git"
|
||||
ARG AITER_BRANCH="v0.1.16.post2"
|
||||
ARG AITER_BRANCH="v0.1.16.post3"
|
||||
ARG AITER_REPO="https://github.com/ROCm/aiter.git"
|
||||
ARG MORI_BRANCH="v1.1.0"
|
||||
ARG MORI_REPO="https://github.com/ROCm/mori.git"
|
||||
|
||||
@@ -249,7 +249,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
|
||||
OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
|
||||
GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \
|
||||
uv pip install -v \
|
||||
uv pip install -v \
|
||||
$ARROW_WHL_FILE \
|
||||
$VISION_WHL_FILE \
|
||||
$HF_XET_WHL_FILE \
|
||||
@@ -257,6 +257,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
$NUMBA_WHL_FILE \
|
||||
$OPENCV_WHL_FILE \
|
||||
$GUIDANCE_WHL_FILE \
|
||||
--torch-backend cpu \
|
||||
--index-strategy unsafe-best-match \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt
|
||||
|
||||
@@ -276,8 +276,9 @@ By default vLLM uses the standard Hugging Face `tokenizers` library to power
|
||||
the fast tokenizer. For BPE tokenizers (Qwen, Llama, DeepSeek, GPT-OSS, etc.)
|
||||
you can switch to the [fastokens](https://github.com/crusoecloud/fastokens)
|
||||
Rust backend, a drop-in replacement that's substantially faster on
|
||||
encode/decode and on streaming detokenization. Enable it by setting
|
||||
`VLLM_USE_FASTOKENS=1`:
|
||||
encode/decode and on streaming detokenization. `VLLM_USE_FASTOKENS` is
|
||||
available in vLLM v0.23.0 and later. If your installed vLLM version does not
|
||||
recognize the environment variable, upgrade vLLM before enabling the override:
|
||||
|
||||
```console
|
||||
VLLM_USE_FASTOKENS=1 vllm serve Qwen/Qwen3-8B
|
||||
|
||||
@@ -168,9 +168,9 @@ Priority is **1 = highest** (tried first).
|
||||
| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ✅ | ❌ | ✅ | All | ≥10.0 |
|
||||
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any |
|
||||
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder Only | Any |
|
||||
| `HPC_ATTN` | | fp16, bf16 | `auto`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
|
||||
| `HPC_ATTN` | | fp16, bf16 | `auto`, `bfloat16`, `fp8_e4m3` | 64 | 128 | ❌ | ❌ | ❌ | ❌ | Decoder | ≥9.0 |
|
||||
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ✅ | ✅ | ❌ | ❌ | Decoder | N/A |
|
||||
| `ROCM_AITER_UNIFIED_ATTN` | | bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
|
||||
| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | ❌ | All | N/A |
|
||||
| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ❌ | ✅ | ✅ | ❌ | Decoder, Encoder, Encoder Only | N/A |
|
||||
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2`, `int4_per_token_head`, `int8_per_token_head`, `fp8_per_token_head` | %16 | Any | ✅ | ✅ | ✅ | ❌ | All | Any |
|
||||
| `TRITON_ATTN_DIFFKV` | | fp16, bf16 | `auto`, `bfloat16` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | Any |
|
||||
|
||||
@@ -816,6 +816,44 @@ Full example: [examples/generate/multimodal/openai_chat_completion_client_for_mu
|
||||
export VLLM_VIDEO_FETCH_TIMEOUT=<timeout>
|
||||
```
|
||||
|
||||
#### Video Decoding Backend
|
||||
|
||||
vLLM decodes video bytes into frames using a selectable decoding backend. Three
|
||||
backends are supported:
|
||||
|
||||
- `opencv` (default): OpenCV-based decoder.
|
||||
- `pyav`: PyAV decoder.
|
||||
- `torchcodec`: TorchCodec (PyTorch-native) decoder.
|
||||
|
||||
All three backends are ultimately backed by FFmpeg. `torchcodec` lets
|
||||
you choose which FFmpeg version is used while `opencv` and `pyav` rely on
|
||||
whichever FFmpeg build they were linked against.
|
||||
|
||||
Select the backend by passing the `backend` parameter via `--media-io-kwargs`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \
|
||||
--media-io-kwargs '{"video": {"backend": "torchcodec"}}'
|
||||
```
|
||||
|
||||
**TorchCodec-specific parameters:**
|
||||
|
||||
The following parameters only apply to the `torchcodec` backend:
|
||||
|
||||
- `num_ffmpeg_threads`: Number of FFmpeg decoding threads. `0` (default) relies
|
||||
on the FFmpeg default, which is `min(cpu_count + 1, 16)`. This allows you to
|
||||
control thread over-subscription.
|
||||
- `seek_mode`: Seek mode for the decoder. `"exact"` (default) guarantees
|
||||
frame-accurate sampling by scanning the file when the decoder is created.
|
||||
`"approximate"` skips that scan for faster decoder creation, at the cost of
|
||||
relying on the file's metadata (which may yield less accurate seeking).
|
||||
|
||||
```bash
|
||||
# Example: TorchCodec with approximate seek mode and 4 FFmpeg threads
|
||||
vllm serve Qwen/Qwen3-VL-30B-A3B-Instruct \
|
||||
--media-io-kwargs '{"video": {"backend": "torchcodec", "seek_mode": "approximate", "num_ffmpeg_threads": 4}}'
|
||||
```
|
||||
|
||||
#### Video Frame Recovery
|
||||
|
||||
For improved robustness when processing potentially corrupted or truncated video files, vLLM supports optional frame recovery using a dynamic window forward-scan approach. When enabled, if a target frame fails to load during sequential reading, the next successfully grabbed frame (before the next target frame) will be used in its place.
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
# Per-Request Metrics
|
||||
|
||||
vLLM can return per-request timing metrics directly in API responses.
|
||||
This is useful for billing, SLA monitoring, and latency analysis at the
|
||||
individual request level, as a complement to the server-aggregated Prometheus
|
||||
metrics exposed at `/metrics`.
|
||||
|
||||
## Enabling
|
||||
|
||||
Start the server with `--enable-per-request-metrics`:
|
||||
|
||||
```bash
|
||||
vllm serve meta-llama/Llama-3.1-8B-Instruct --enable-per-request-metrics
|
||||
```
|
||||
|
||||
When this flag is set, supported API responses include metrics for each
|
||||
attributable request.
|
||||
|
||||
!!! note
|
||||
At high concurrency, enabling per-request metrics computation may introduce
|
||||
non-negligible CPU overhead. Benchmark your specific workload to evaluate the
|
||||
impact before enabling in production.
|
||||
|
||||
## Response Format
|
||||
|
||||
When per-request metrics are enabled, the response includes a `metrics` object:
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "chatcmpl-abc123",
|
||||
"object": "chat.completion",
|
||||
"model": "meta-llama/Llama-3.1-8B-Instruct",
|
||||
"choices": [ ... ],
|
||||
"usage": {
|
||||
"prompt_tokens": 42,
|
||||
"completion_tokens": 128,
|
||||
"total_tokens": 170
|
||||
},
|
||||
"metrics": {
|
||||
"time_to_first_token_ms": 85.2,
|
||||
"generation_time_ms": 1240.5,
|
||||
"queue_time_ms": 12.3,
|
||||
"mean_itl_ms": 9.1,
|
||||
"tokens_per_second": 103.2
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Description |
|
||||
| --- | --- |
|
||||
| `time_to_first_token_ms` | Time from when the request was scheduled until the first output token was generated (TTFT). |
|
||||
| `generation_time_ms` | Decode time: time from the first output token to the last output token. Excludes both queue wait and prefill/TTFT. |
|
||||
| `queue_time_ms` | Time the request spent waiting in the scheduler queue before processing began. |
|
||||
| `mean_itl_ms` | Mean inter-token latency (average time between successive output tokens) during the decode phase. `null` for single-token responses. |
|
||||
| `tokens_per_second` | Overall output token throughput: all generated tokens over the inference interval (scheduling to last output token). Unlike `generation_time_ms`, this includes the prefill phase, so it reflects end-to-end generation speed rather than pure decode speed. |
|
||||
|
||||
All fields are `null` if the underlying timing data is not available for that
|
||||
request.
|
||||
|
||||
!!! note
|
||||
Timing metrics describe a single generation stream, so they are only
|
||||
returned when the request maps to exactly one. They are suppressed (the
|
||||
`metrics` object is `null`) for requests with `n > 1`, because the
|
||||
underlying timing data reflects only one of the `n` sequences and cannot be
|
||||
accurately attributed to the request as a whole. Token usage
|
||||
(`prompt_tokens`, `completion_tokens`) remains accurate in these cases.
|
||||
Per-request metrics also require server-side statistics logging, which is
|
||||
on by default. vLLM rejects `--enable-per-request-metrics` when
|
||||
`--disable-log-stats` is also set.
|
||||
|
||||
## Example Request
|
||||
|
||||
=== "Non-streaming"
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8000/v1", api_key="token")
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="meta-llama/Llama-3.1-8B-Instruct",
|
||||
messages=[{"role": "user", "content": "What is the capital of France?"}],
|
||||
)
|
||||
|
||||
print(response.usage)
|
||||
print(response.model_extra.get("metrics"))
|
||||
```
|
||||
|
||||
=== "Streaming"
|
||||
|
||||
In streaming responses, metrics are attached to the final usage chunk (the
|
||||
chunk sent after all content chunks). That chunk is only emitted when usage
|
||||
reporting is enabled with `stream_options.include_usage: true` or forced
|
||||
server-side with `--enable-force-include-usage`. Without forced usage, a
|
||||
streaming client must set `stream_options.include_usage: true` to receive
|
||||
metrics.
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8000/v1", api_key="token")
|
||||
|
||||
stream = client.chat.completions.create(
|
||||
model="meta-llama/Llama-3.1-8B-Instruct",
|
||||
messages=[{"role": "user", "content": "What is the capital of France?"}],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
)
|
||||
|
||||
for chunk in stream:
|
||||
if chunk.usage:
|
||||
print("Usage:", chunk.usage)
|
||||
print("Metrics:", chunk.model_extra.get("metrics"))
|
||||
```
|
||||
|
||||
## Completions API
|
||||
|
||||
Per-request metrics are also available on the `/v1/completions` endpoint using
|
||||
the same `metrics` response field. As with `n > 1`, metrics are omitted for
|
||||
requests with multiple prompts, because the timing data cannot be attributed to
|
||||
a single prompt's generation.
|
||||
|
||||
## Relationship to Prometheus Metrics
|
||||
|
||||
The `metrics` response field provides per-request values for a single request.
|
||||
The `/metrics` Prometheus endpoint exposes server-level histograms (e.g.
|
||||
`vllm:time_to_first_token_seconds`) that aggregate across all requests.
|
||||
@@ -75,14 +75,11 @@ vllm serve Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound \
|
||||
--max-model-len 4096
|
||||
```
|
||||
|
||||
!!! note
|
||||
To deploy `wNa16` models on Intel GPU/CPU, please add `--enforce-eager` for now.
|
||||
|
||||
## Evaluating the Quantized Model with vLLM
|
||||
|
||||
```bash
|
||||
lm_eval --model vllm \
|
||||
--model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048,enforce_eager=True" \
|
||||
--model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048" \
|
||||
--tasks gsm8k \
|
||||
--num_fewshot 5 \
|
||||
--batch_size 128
|
||||
|
||||
@@ -439,7 +439,7 @@ Additionally, to enable structured output, you'll need to create a new `Reasoner
|
||||
end_token: str = "</think>"
|
||||
|
||||
@classmethod
|
||||
def from_tokenizer(cls, tokenizer: PreTrainedTokenizer) -> Reasoner:
|
||||
def from_tokenizer(cls, tokenizer: PythonBackend) -> Reasoner:
|
||||
return cls(
|
||||
start_token_id=tokenizer.encode("<think>", add_special_tokens=False)[0],
|
||||
end_token_id=tokenizer.encode("</think>", add_special_tokens=False)[0],
|
||||
|
||||
@@ -73,3 +73,4 @@ VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \
|
||||
|
||||
* Tested with Eagle, Eagle-3, and DFlash. Other SD methods may or may not work out of the box
|
||||
* Full Cudagraph only works with Model Runner V2. MRv1 only supports piece-wise cuda graph with this feature
|
||||
* Not compatible with data parallelism (`--data-parallel-size > 1`). Each DP rank schedules independently, so ranks can pick different K values, causing DP collective divergence and deadlocks. When DP is enabled, vLLM automatically disables `num_speculative_tokens_per_batch_size` and falls back to the static `num_speculative_tokens` value.
|
||||
|
||||
@@ -35,15 +35,10 @@ After installation of XCode and the Command Line Tools, which include Apple Clan
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
uv pip install -r requirements/cpu.txt --index-strategy unsafe-best-match
|
||||
uv pip install -r requirements/cpu.txt
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
!!! tip
|
||||
The `--index-strategy unsafe-best-match` flag is needed to resolve dependencies across multiple package indexes (PyTorch CPU index and PyPI). Without this flag, you may encounter `typing-extensions` version conflicts.
|
||||
|
||||
The term "unsafe" refers to the package resolution strategy, not security. By default, `uv` only searches the first index where a package is found to prevent dependency confusion attacks. This flag allows `uv` to search all configured indexes to find the best compatible versions. Since both PyTorch and PyPI are trusted package sources, using this strategy is safe and appropriate for vLLM installation.
|
||||
|
||||
!!! note
|
||||
On macOS the `VLLM_TARGET_DEVICE` is automatically set to `cpu`, which is currently the only supported device.
|
||||
|
||||
|
||||
@@ -48,10 +48,10 @@ Execute the following commands to build and install vLLM from source.
|
||||
|
||||
```bash
|
||||
uv pip install -v \
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu \
|
||||
--torch-backend auto \
|
||||
-r requirements/build/cpu.txt \
|
||||
-r requirements/cpu.txt \
|
||||
--torch-backend cpu \
|
||||
--index-strategy unsafe-best-match && \
|
||||
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
|
||||
uv pip install dist/*.whl
|
||||
```
|
||||
|
||||
@@ -2,8 +2,7 @@
|
||||
|
||||
!!! note
|
||||
We currently support pooling models primarily for convenience. This is not guaranteed to provide any performance
|
||||
improvements over using Hugging Face Transformers or Sentence Transformers directly.
|
||||
|
||||
improvements over using Hugging Face Transformers or Sentence Transformers directly.
|
||||
We plan to optimize pooling models in vLLM. Please comment on <https://github.com/vllm-project/vllm/issues/21796> if you have any suggestions!
|
||||
|
||||
## What are pooling models?
|
||||
@@ -63,7 +62,7 @@ please refer to [IO Processor Plugins](../../design/io_processor_plugins.md).
|
||||
|
||||
!!! note
|
||||
Within classification tasks, there is a specialized subcategory: Cross-encoder (aka reranker) models. These models
|
||||
are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
|
||||
are a subset of classification models that accept two prompts as input and output num_labels equal to 1.
|
||||
|
||||
### Pooling Types
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ These models are what we list in [supported text models](#list-of-text-only-lang
|
||||
|
||||
### Transformers
|
||||
|
||||
vLLM also supports model implementations that are available in Transformers. You should expect the performance of a Transformers model implementation used in vLLM to be within <5% of the performance of a dedicated vLLM model implementation. We call this feature the "Transformers modeling backend".
|
||||
vLLM also supports model implementations that are available in Transformers. We call this feature the "Transformers modeling backend". The performance of models loaded with the Transformers modeling backend should be identical to a dedicated vLLM model implementation.
|
||||
|
||||
Currently, the Transformers modeling backend works for the following:
|
||||
|
||||
@@ -140,7 +140,7 @@ Here is what happens in the background when this model is loaded:
|
||||
|
||||
That's it!
|
||||
|
||||
For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you must add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
|
||||
For your model to be compatible with vLLM's tensor parallel and/or pipeline parallel features, you may need to add `base_model_tp_plan` and/or `base_model_pp_plan` to your model's config class:
|
||||
|
||||
<details class="code">
|
||||
<summary>configuration_my_model.py</summary>
|
||||
@@ -168,9 +168,11 @@ class MyConfig(PretrainedConfig):
|
||||
</details>
|
||||
|
||||
- `base_model_tp_plan` is a `dict` that maps fully qualified layer name patterns to tensor parallel styles (currently only `"colwise"` and `"rowwise"` are supported).
|
||||
- vLLM infers the tensor parallel style of standard attention (`q`/`k`/`v`/`o_proj`) and gated-MLP/experts (`gate`/`up`/`down_proj`) projections if it can fuse them, so these may not need to be listed. `base_model_tp_plan` is only _required_ for layers that do not follow these patterns; any linear that is neither fused nor named in the plan is replicated.
|
||||
- `base_model_pp_plan` is a `dict` that maps direct child layer names to `tuple`s of `list`s of `str`s:
|
||||
- You only need to do this for layers which are not present on all pipeline stages
|
||||
- vLLM assumes that there will be only one `nn.ModuleList`, which is distributed across the pipeline stages
|
||||
- When no `base_model_pp_plan` is provided, the Transformers modelling backend infers the split from the text model's sole `nn.ModuleList`, keeping the parameter-bearing modules around it (input embeddings, final norm) on the first/last stage (depending on declaration order) and parameter-free modules (e.g. rotary embeddings) on every stage
|
||||
- The `list` in the first element of the `tuple` contains the names of the input arguments
|
||||
- The `list` in the last element of the `tuple` contains the names of the variables the layer outputs to in your modeling code
|
||||
|
||||
@@ -240,50 +242,24 @@ Use the Hugging Face CLI to [manage models](https://huggingface.co/docs/huggingf
|
||||
|
||||
```bash
|
||||
# List cached models
|
||||
hf scan-cache
|
||||
hf cache list -q
|
||||
|
||||
# Show detailed (verbose) output
|
||||
hf scan-cache -v
|
||||
hf cache list
|
||||
|
||||
# Specify a custom cache directory
|
||||
hf scan-cache --dir ~/.cache/huggingface/hub
|
||||
hf cache list --dir ~/.cache/huggingface/hub
|
||||
```
|
||||
|
||||
#### Delete a cached model
|
||||
|
||||
Use the Hugging Face CLI to interactively [delete downloaded model](https://huggingface.co/docs/huggingface_hub/guides/manage-cache#clean-your-cache) from the cache:
|
||||
Use the Hugging Face CLI to [delete downloaded model](https://huggingface.co/docs/huggingface_hub/guides/manage-cache#clean-your-cache) from the cache:
|
||||
|
||||
<details>
|
||||
<summary>Commands</summary>
|
||||
|
||||
```console
|
||||
# The `delete-cache` command requires extra dependencies to work with the TUI.
|
||||
# Please run `pip install huggingface_hub[cli]` to install them.
|
||||
|
||||
# Launch the interactive TUI to select models to delete
|
||||
$ hf delete-cache
|
||||
? Select revisions to delete: 1 revisions selected counting for 438.9M.
|
||||
○ None of the following (if selected, nothing will be deleted).
|
||||
Model BAAI/bge-base-en-v1.5 (438.9M, used 1 week ago)
|
||||
❯ ◉ a5beb1e3: main # modified 1 week ago
|
||||
|
||||
Model BAAI/bge-large-en-v1.5 (1.3G, used 1 week ago)
|
||||
○ d4aa6901: main # modified 1 week ago
|
||||
|
||||
Model BAAI/bge-reranker-base (1.1G, used 4 weeks ago)
|
||||
○ 2cfc18c9: main # modified 4 weeks ago
|
||||
|
||||
Press <space> to select, <enter> to validate and <ctrl+c> to quit without modification.
|
||||
|
||||
# Need to confirm after selected
|
||||
? Select revisions to delete: 1 revision(s) selected.
|
||||
? 1 revisions selected counting for 438.9M. Confirm deletion ? Yes
|
||||
Start deletion.
|
||||
Done. Deleted 1 repo(s) and 0 revision(s) for a total of 438.9M.
|
||||
```bash
|
||||
# delete all the cached objects
|
||||
hf cache rm $(hf cache list -q)
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
#### Using a proxy
|
||||
|
||||
Here are some tips for loading/downloading models from Hugging Face using a proxy:
|
||||
|
||||
@@ -19,7 +19,7 @@ Run:
|
||||
"""
|
||||
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizer
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, PythonBackend
|
||||
|
||||
from vllm import LLM
|
||||
|
||||
@@ -34,7 +34,7 @@ def init_tokenizer_and_llm(model_name: str):
|
||||
|
||||
def get_prompt_embeds(
|
||||
chat: list[dict[str, str]],
|
||||
tokenizer: PreTrainedTokenizer,
|
||||
tokenizer: PythonBackend,
|
||||
embedding_layer: torch.nn.Module,
|
||||
):
|
||||
token_ids = tokenizer.apply_chat_template(
|
||||
@@ -45,7 +45,7 @@ def get_prompt_embeds(
|
||||
|
||||
|
||||
def single_prompt_inference(
|
||||
llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module
|
||||
llm: LLM, tokenizer: PythonBackend, embedding_layer: torch.nn.Module
|
||||
):
|
||||
chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
|
||||
prompt_embeds = get_prompt_embeds(chat, tokenizer, embedding_layer)
|
||||
@@ -64,7 +64,7 @@ def single_prompt_inference(
|
||||
|
||||
|
||||
def batch_prompt_inference(
|
||||
llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module
|
||||
llm: LLM, tokenizer: PythonBackend, embedding_layer: torch.nn.Module
|
||||
):
|
||||
chats = [
|
||||
[{"role": "user", "content": "Please tell me about the capital of France."}],
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# ruff: noqa: E501
|
||||
|
||||
"""
|
||||
Example online usage of the Jina Reranker v3 score and rerank APIs with a task
|
||||
instruction.
|
||||
|
||||
Run `vllm serve jinaai/jina-reranker-v3 --runner pooling` to start up the
|
||||
server in vLLM.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
|
||||
headers = {"User-Agent": "Test Client"}
|
||||
response = requests.post(api_url, headers=headers, json=prompt)
|
||||
return response
|
||||
|
||||
|
||||
def print_response(name: str, prompt: dict, response: requests.Response) -> None:
|
||||
print(f"\n{name} request:")
|
||||
print(json.dumps(prompt, indent=2))
|
||||
print(f"\n{name} response:")
|
||||
print(json.dumps(response.json(), indent=2))
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--model", type=str, default="jinaai/jina-reranker-v3")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main(args):
|
||||
score_url = f"http://{args.host}:{args.port}/score"
|
||||
rerank_url = f"http://{args.host}:{args.port}/rerank"
|
||||
model_name = args.model
|
||||
|
||||
query = "Which passage is about sports?"
|
||||
documents = [
|
||||
"Basketball is played by two teams on a court.",
|
||||
"Green tea contains antioxidants and may support metabolism.",
|
||||
]
|
||||
instruction = "Rank passages about sports higher than passages about nutrition."
|
||||
|
||||
score_prompt = {
|
||||
"model": model_name,
|
||||
"queries": query,
|
||||
"documents": documents,
|
||||
"instruction": instruction,
|
||||
}
|
||||
score_response = post_http_request(prompt=score_prompt, api_url=score_url)
|
||||
print_response("Score", score_prompt, score_response)
|
||||
|
||||
rerank_prompt = {
|
||||
"model": model_name,
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
"instruction": instruction,
|
||||
}
|
||||
rerank_response = post_http_request(prompt=rerank_prompt, api_url=rerank_url)
|
||||
print_response("Rerank", rerank_prompt, rerank_response)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
@@ -1,4 +1,3 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu
|
||||
cmake>=3.26.1
|
||||
ninja
|
||||
packaging>=24.2
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/cpu
|
||||
# Common dependencies
|
||||
-r common.txt
|
||||
|
||||
@@ -16,6 +15,9 @@ torchaudio; platform_machine != "s390x" and platform_machine != "riscv64"
|
||||
# required for the image processor of phi3v, this must be updated alongside torch
|
||||
torchvision; platform_machine != "s390x" and platform_machine != "riscv64"
|
||||
|
||||
# required for the torchcodec video decoding backend
|
||||
torchcodec >= 0.14; platform_machine != "s390x" and platform_machine != "riscv64" and platform_machine != "ppc64le"
|
||||
|
||||
# Intel Extension for PyTorch, only for x86_64 CPUs
|
||||
intel-openmp==2024.2.1; platform_machine == "x86_64"
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ torch==2.11.0
|
||||
torchaudio==2.11.0
|
||||
# These must be updated alongside torch
|
||||
torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
|
||||
torchcodec >= 0.14
|
||||
PyNvVideoCodec==2.1.0
|
||||
# FlashInfer should be updated together with the Dockerfile
|
||||
flashinfer-python==0.6.13
|
||||
@@ -25,7 +26,7 @@ nvidia-cutlass-dsl[cu13]==4.5.2
|
||||
quack-kernels>=0.3.3
|
||||
|
||||
# Tokenspeed_MLA for faster mla with spec decode
|
||||
tokenspeed-mla==0.1.2
|
||||
tokenspeed-mla==0.1.2; platform_system == "Linux"
|
||||
|
||||
# Humming kernels for quantization gemm
|
||||
humming-kernels[cu13]==0.1.6
|
||||
humming-kernels[cu13]==0.1.10
|
||||
|
||||
@@ -672,7 +672,7 @@ pathvalidate==3.2.1
|
||||
# via pytablewriter
|
||||
patsy==1.0.1
|
||||
# via statsmodels
|
||||
peft==0.18.1
|
||||
peft==0.19.1
|
||||
# via -r requirements/test/cuda.in
|
||||
perceptron==0.1.4
|
||||
# via -r requirements/test/cuda.in
|
||||
@@ -1134,6 +1134,8 @@ torchaudio==2.11.0+cpu
|
||||
# -r requirements/test/cuda.in
|
||||
# encodec
|
||||
# vocos
|
||||
torchcodec==0.14.0+cpu
|
||||
# via -r requirements/test/cuda.in
|
||||
torchvision==0.26.0+cpu
|
||||
# via
|
||||
# -r requirements/test/cuda.in
|
||||
@@ -1157,7 +1159,7 @@ tqdm==4.67.3
|
||||
# segmentation-models-pytorch
|
||||
# sentence-transformers
|
||||
# transformers
|
||||
transformers==5.5.3
|
||||
transformers==5.10.4
|
||||
# via
|
||||
# -r requirements/test/../common.txt
|
||||
# -r requirements/test/cuda.in
|
||||
|
||||
@@ -13,13 +13,14 @@ pytest-cov
|
||||
# testing utils
|
||||
albumentations # required for Nemotron Parse in test_common.py
|
||||
av # required for audio_in_video tests
|
||||
torchcodec >= 0.14 # required for torchcodec video backend tests
|
||||
backoff # required for phi4mm test
|
||||
blobfile # required for kimi-vl test
|
||||
httpx
|
||||
librosa # required for audio tests
|
||||
vector_quantize_pytorch # required for minicpmo_26 test
|
||||
vocos # required for minicpmo_26 test
|
||||
peft>=0.18.1 # required for phi-4-mm test
|
||||
peft>=0.19.1 # required for phi-4-mm test
|
||||
pqdm
|
||||
ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
|
||||
sentence-transformers>=5.2.0 # required for embedding tests
|
||||
@@ -38,7 +39,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_
|
||||
datamodel_code_generator # required for minicpm3 test
|
||||
lm-eval[api]>=0.4.12 # required for model evaluation test
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
transformers==5.5.3
|
||||
transformers==5.10.4
|
||||
tokenizers==0.22.2
|
||||
schemathesis>=4.0.0 # Required for openai schema test.
|
||||
# quantization
|
||||
|
||||
@@ -755,7 +755,7 @@ pathvalidate==3.2.1
|
||||
# via pytablewriter
|
||||
patsy==1.0.1
|
||||
# via statsmodels
|
||||
peft==0.18.1
|
||||
peft==0.19.1
|
||||
# via -r requirements/test/cuda.in
|
||||
perceptron==0.1.4
|
||||
# via -r requirements/test/cuda.in
|
||||
@@ -1233,6 +1233,10 @@ torchaudio==2.11.0+cu130
|
||||
# -r requirements/test/cuda.in
|
||||
# encodec
|
||||
# vocos
|
||||
torchcodec==0.14.0+cu130
|
||||
# via
|
||||
# -c requirements/cuda.txt
|
||||
# -r requirements/test/cuda.in
|
||||
torchvision==0.26.0+cu130
|
||||
# via
|
||||
# -c requirements/cuda.txt
|
||||
@@ -1257,7 +1261,7 @@ tqdm==4.67.3
|
||||
# segmentation-models-pytorch
|
||||
# sentence-transformers
|
||||
# transformers
|
||||
transformers==5.5.3
|
||||
transformers==5.10.4
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
|
||||
@@ -29,7 +29,7 @@ opencv-python-headless >= 4.13.0 # required for video test
|
||||
datamodel_code_generator # required for minicpm3 test
|
||||
lm-eval[api]>=0.4.12 # required for model evaluation test
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
transformers==5.5.3
|
||||
transformers==5.10.4
|
||||
tokenizers==0.22.2
|
||||
schemathesis>=4.0.0 # Required for openai schema test.
|
||||
# quantization
|
||||
|
||||
@@ -19,7 +19,7 @@ httpx
|
||||
librosa # required for audio tests
|
||||
vector_quantize_pytorch # required for minicpmo_26 test
|
||||
vocos # required for minicpmo_26 test
|
||||
peft>=0.15.0 # required for phi-4-mm test
|
||||
peft>=0.19.1 # required for phi-4-mm test
|
||||
pqdm
|
||||
ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline parallelism tests
|
||||
sentence-transformers>=5.2.0 # required for embedding tests
|
||||
@@ -35,7 +35,7 @@ open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_
|
||||
datamodel_code_generator # required for minicpm3 test
|
||||
lm-eval[api]>=0.4.12 # required for model evaluation test
|
||||
mteb[bm25s]>=2, <3 # required for mteb test
|
||||
transformers==5.5.3
|
||||
transformers==5.10.4
|
||||
tokenizers==0.22.2
|
||||
schemathesis>=4.0.0 # Required for openai schema test
|
||||
# quantization
|
||||
|
||||
@@ -731,7 +731,7 @@ pathvalidate==3.3.1
|
||||
# via pytablewriter
|
||||
patsy==1.0.2
|
||||
# via statsmodels
|
||||
peft==0.18.1
|
||||
peft==0.19.1
|
||||
# via -r requirements/test/rocm.in
|
||||
perceptron==0.1.4
|
||||
# via -r requirements/test/rocm.in
|
||||
@@ -1218,7 +1218,7 @@ tqdm==4.67.3
|
||||
# sentence-transformers
|
||||
# tilelang
|
||||
# transformers
|
||||
transformers==5.5.3
|
||||
transformers==5.10.4
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
|
||||
@@ -17,6 +17,7 @@ accelerate
|
||||
arctic-inference
|
||||
lm_eval[api]>=0.4.12
|
||||
modelscope<1.38
|
||||
transformers==5.10.4
|
||||
|
||||
# --- Audio Processing ---
|
||||
librosa
|
||||
|
||||
@@ -940,10 +940,11 @@ tqdm==4.67.3
|
||||
# pqdm
|
||||
# sentence-transformers
|
||||
# transformers
|
||||
transformers==5.5.3
|
||||
transformers==5.10.4
|
||||
# via
|
||||
# -c requirements/common.txt
|
||||
# -r requirements/test/../common.txt
|
||||
# -r requirements/test/xpu.in
|
||||
# compressed-tensors
|
||||
# sentence-transformers
|
||||
# xgrammar
|
||||
|
||||
@@ -12,4 +12,4 @@ ray[data]
|
||||
setuptools==78.1.0
|
||||
setuptools-rust>=1.9.0
|
||||
nixl==0.3.0
|
||||
tpu-inference==0.23.0
|
||||
tpu-inference==0.24.0
|
||||
|
||||
@@ -15,6 +15,7 @@ numba == 0.65.0 # Required for N-gram speculative decoding
|
||||
torch==2.12.0
|
||||
torchaudio
|
||||
torchvision
|
||||
torchcodec >= 0.14 # Required for the torchcodec video decoding backend
|
||||
|
||||
auto_round_lib>=0.13.3
|
||||
auto_round_lib>=0.14.0
|
||||
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10.1/vllm_xpu_kernels-0.1.10.1-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
|
||||
@@ -172,6 +172,9 @@ impl ChatLlm {
|
||||
pub async fn chat(&self, mut request: ChatRequest) -> Result<ChatEventStream> {
|
||||
request.validate()?;
|
||||
|
||||
// Stamp before rendering so render and tokenize count toward TTFT/e2e.
|
||||
let arrival_time = vllm_llm::current_unix_timestamp_secs();
|
||||
|
||||
let output_processor = self.backend.new_chat_output_processor(
|
||||
&mut request,
|
||||
NewChatOutputProcessorOptions {
|
||||
@@ -210,6 +213,7 @@ impl ChatLlm {
|
||||
data_parallel_rank: request.data_parallel_rank,
|
||||
reasoning_parser_kwargs,
|
||||
lora_request: request.lora_request,
|
||||
arrival_time: Some(arrival_time),
|
||||
};
|
||||
let decoded_stream = self.text.generate(text_request).await?.map_err(Error::from).boxed();
|
||||
|
||||
|
||||
@@ -130,6 +130,19 @@ impl RoundtripCase {
|
||||
}
|
||||
}
|
||||
|
||||
/// DeepSeek V3.2 DSML tool-call format.
|
||||
fn deepseek_v32() -> Self {
|
||||
Self {
|
||||
model_id: "deepseek-ai/DeepSeek-V3.2-Exp",
|
||||
assistant_stop_suffix: "<|end▁of▁sentence|>",
|
||||
tool_call_parser: ParserSelection::Auto,
|
||||
reasoning_parser: ParserSelection::Auto,
|
||||
thinking_behavior: ThinkingBehavior::Toggleable { default: false },
|
||||
json_fmt: compact_json_fmt(),
|
||||
sort_json_keys: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// GLM-4.7 XML-like argument format with `<think>` reasoning tags.
|
||||
fn glm47() -> Self {
|
||||
Self {
|
||||
@@ -235,6 +248,7 @@ roundtrip_tests! {
|
||||
qwen35 => [reasoning_and_content, tool_call_mix],
|
||||
minimax_m25 => [reasoning_and_content, tool_call_mix],
|
||||
deepseek_v4 => [reasoning_and_content, tool_call_mix],
|
||||
deepseek_v32 => [tool_call_mix],
|
||||
glm47 => [reasoning_and_content, tool_call_mix],
|
||||
seed_oss => [reasoning_and_content],
|
||||
step3p5 => [reasoning_and_content],
|
||||
|
||||
@@ -456,13 +456,17 @@ pub struct ServerUnsupportedArgs {
|
||||
|
||||
/// Enable the `/tokenizer_info` endpoint. May expose chat
|
||||
/// templates and other tokenizer configuration.
|
||||
///
|
||||
/// Accepted as a no-op: the Rust frontend serves `/tokenize` and
|
||||
/// `/detokenize`, but does not implement `/tokenizer_info` yet.
|
||||
#[arg(
|
||||
long,
|
||||
visible_alias = "no-enable-tokenizer-info-endpoint",
|
||||
default_missing_value = "true",
|
||||
num_args = 0..=1
|
||||
num_args = 0..=1,
|
||||
hide = true
|
||||
)]
|
||||
pub enable_tokenizer_info_endpoint: Option<Unsupported>,
|
||||
pub enable_tokenizer_info_endpoint: Option<Noop>,
|
||||
|
||||
/// If set to True, log model outputs (generations).
|
||||
/// Requires `--enable-log-requests`. As with `--enable-log-requests`,
|
||||
|
||||
@@ -512,6 +512,7 @@ impl EngineCoreClient {
|
||||
|
||||
Ok(EngineCoreOutputStream::new(
|
||||
request_id,
|
||||
engine_id.engine_index().unwrap_or(0),
|
||||
self.abort_tx.clone(),
|
||||
rx,
|
||||
))
|
||||
|
||||
@@ -15,7 +15,7 @@ use crate::client::state::{OutputReceiver, RequestRegistry, UtilityReceiver, Uti
|
||||
use crate::client::stream::EngineCoreStreamOutput;
|
||||
use crate::client::{AbortCause, AbortRequest};
|
||||
use crate::error::{client_closed, dispatcher_closed, unexpected_dispatcher_output};
|
||||
use crate::metrics::{LoraInfoExporter, record_scheduler_stats};
|
||||
use crate::metrics::{LoraInfoExporter, SchedulerStatsRecorder};
|
||||
use crate::protocol::encode_msgpack;
|
||||
use crate::protocol::output::{EngineCoreOutput, EngineCoreOutputs};
|
||||
use crate::protocol::request::EngineCoreRequestType;
|
||||
@@ -29,6 +29,7 @@ pub(crate) struct ClientInner {
|
||||
/// The runtime handle used for sending messages to the engine.
|
||||
handle: Handle,
|
||||
model_name: String,
|
||||
scheduler_stats_recorder: SchedulerStatsRecorder,
|
||||
request_reg: Mutex<RequestRegistry>,
|
||||
utility_reg: Mutex<UtilityRegistry>,
|
||||
health_error: ArcSwapOption<Error>,
|
||||
@@ -43,10 +44,13 @@ impl ClientInner {
|
||||
model_name: String,
|
||||
engines: &[ConnectedEngine],
|
||||
) -> Self {
|
||||
let scheduler_stats_recorder =
|
||||
SchedulerStatsRecorder::new(&METRICS.scheduler, &model_name, engines);
|
||||
Self {
|
||||
input_send,
|
||||
handle,
|
||||
model_name,
|
||||
scheduler_stats_recorder,
|
||||
request_reg: Mutex::new(RequestRegistry::new(engines)),
|
||||
utility_reg: Mutex::new(UtilityRegistry::default()),
|
||||
health_error: ArcSwapOption::empty(),
|
||||
@@ -389,12 +393,7 @@ pub(crate) async fn run_output_dispatcher_loop(
|
||||
"dropping scheduler stats for unknown engine"
|
||||
);
|
||||
}
|
||||
record_scheduler_stats(
|
||||
&METRICS.scheduler,
|
||||
inner.model_name(),
|
||||
batch.engine_index,
|
||||
scheduler_stats,
|
||||
);
|
||||
inner.scheduler_stats_recorder.record(batch.engine_index, scheduler_stats);
|
||||
}
|
||||
|
||||
// The engine's scheduler stats never carry adapter names;
|
||||
|
||||
@@ -45,6 +45,7 @@ impl Deref for EngineCoreStreamOutput {
|
||||
/// `finish_reason` is non-`None`.
|
||||
pub struct EngineCoreOutputStream {
|
||||
request_id: String,
|
||||
engine_index: u32,
|
||||
abort_tx: mpsc::UnboundedSender<AbortRequest>,
|
||||
state: State,
|
||||
rx: OutputReceiver,
|
||||
@@ -53,11 +54,13 @@ pub struct EngineCoreOutputStream {
|
||||
impl EngineCoreOutputStream {
|
||||
pub(crate) fn new(
|
||||
request_id: String,
|
||||
engine_index: u32,
|
||||
abort_tx: mpsc::UnboundedSender<AbortRequest>,
|
||||
rx: OutputReceiver,
|
||||
) -> Self {
|
||||
Self {
|
||||
request_id,
|
||||
engine_index,
|
||||
abort_tx,
|
||||
state: State::Running,
|
||||
rx,
|
||||
@@ -68,6 +71,11 @@ impl EngineCoreOutputStream {
|
||||
pub fn request_id(&self) -> &str {
|
||||
&self.request_id
|
||||
}
|
||||
|
||||
/// Return the index of the engine that owns this request.
|
||||
pub fn engine_index(&self) -> u32 {
|
||||
self.engine_index
|
||||
}
|
||||
}
|
||||
|
||||
impl Stream for EngineCoreOutputStream {
|
||||
|
||||
@@ -1,90 +1,190 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::collections::BTreeSet;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use vllm_metrics::{
|
||||
EngineLabels, EnginePositionLabels, LoraAdapterNames, LoraInfoLabels, SchedulerMetrics,
|
||||
EngineLabels, EnginePositionLabels, F64Gauge, Family, HistogramMetric, LoraAdapterNames,
|
||||
LoraInfoLabels, SchedulerLogStatsAccumulator, SchedulerMetrics, U64Counter, U64Gauge,
|
||||
WaitingReasonLabels,
|
||||
};
|
||||
|
||||
use crate::protocol::stats::SchedulerStats;
|
||||
use crate::transport::ConnectedEngine;
|
||||
|
||||
const WAITING_REASON_CAPACITY: &str = "capacity";
|
||||
const WAITING_REASON_DEFERRED: &str = "deferred";
|
||||
|
||||
/// Record the scheduler-stats-backed metrics for one engine at one point in
|
||||
/// time.
|
||||
pub(crate) fn record_scheduler_stats(
|
||||
metrics: &SchedulerMetrics,
|
||||
model_name: impl Into<String>,
|
||||
engine: u32,
|
||||
stats: &SchedulerStats,
|
||||
) {
|
||||
let model_name = model_name.into();
|
||||
let labels = EngineLabels {
|
||||
model_name: model_name.clone(),
|
||||
engine,
|
||||
};
|
||||
/// Cached scheduler-stats metric handles for all engines connected to one
|
||||
/// frontend client.
|
||||
pub(crate) struct SchedulerStatsRecorder {
|
||||
engines: BTreeMap<u32, SchedulerStatsHandles>,
|
||||
}
|
||||
|
||||
/// Per-engine cached metric handles used while recording `SchedulerStats`.
|
||||
struct SchedulerStatsHandles {
|
||||
// Base labels reused for dynamic child labels.
|
||||
labels: EngineLabels,
|
||||
|
||||
// Scheduler state gauges.
|
||||
metrics.scheduler_running.get_or_create(&labels).set(stats.num_running_reqs);
|
||||
metrics
|
||||
.scheduler_waiting
|
||||
.get_or_create(&labels)
|
||||
.set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs);
|
||||
metrics
|
||||
.scheduler_waiting_by_reason
|
||||
.get_or_create(&WaitingReasonLabels {
|
||||
model_name: model_name.clone(),
|
||||
engine,
|
||||
reason: WAITING_REASON_CAPACITY,
|
||||
})
|
||||
.set(stats.num_waiting_reqs);
|
||||
metrics
|
||||
.scheduler_waiting_by_reason
|
||||
.get_or_create(&WaitingReasonLabels {
|
||||
model_name: model_name.clone(),
|
||||
engine,
|
||||
reason: WAITING_REASON_DEFERRED,
|
||||
})
|
||||
.set(stats.num_skipped_waiting_reqs);
|
||||
metrics.kv_cache_usage.get_or_create(&labels).set(stats.kv_cache_usage);
|
||||
scheduler_running: U64Gauge,
|
||||
scheduler_waiting: U64Gauge,
|
||||
scheduler_waiting_capacity: U64Gauge,
|
||||
scheduler_waiting_deferred: U64Gauge,
|
||||
kv_cache_usage: F64Gauge,
|
||||
|
||||
// Prefix-cache counters, including the connector-backed external cache path.
|
||||
metrics
|
||||
.prefix_cache_queries
|
||||
.get_or_create(&labels)
|
||||
.inc_by(stats.prefix_cache_stats.base.queries);
|
||||
metrics
|
||||
.prefix_cache_hits
|
||||
.get_or_create(&labels)
|
||||
.inc_by(stats.prefix_cache_stats.base.hits);
|
||||
prefix_cache_queries: U64Counter,
|
||||
prefix_cache_hits: U64Counter,
|
||||
external_prefix_cache_queries: U64Counter,
|
||||
external_prefix_cache_hits: U64Counter,
|
||||
|
||||
// Speculative decoding counters.
|
||||
spec_decode_num_drafts: U64Counter,
|
||||
spec_decode_num_draft_tokens: U64Counter,
|
||||
spec_decode_num_accepted_tokens: U64Counter,
|
||||
spec_decode_num_accepted_tokens_per_pos: Family<EnginePositionLabels, U64Counter>,
|
||||
|
||||
// Per-engine performance / MFU counters.
|
||||
estimated_flops_per_gpu: U64Counter,
|
||||
estimated_read_bytes_per_gpu: U64Counter,
|
||||
estimated_write_bytes_per_gpu: U64Counter,
|
||||
|
||||
// Sampled KV-cache residency histograms.
|
||||
kv_block_lifetime_seconds: HistogramMetric,
|
||||
kv_block_idle_before_evict_seconds: HistogramMetric,
|
||||
kv_block_reuse_gap_seconds: HistogramMetric,
|
||||
|
||||
// Non-Prometheus interval accumulator for periodic text-log helpers.
|
||||
log_stats: SchedulerLogStatsAccumulator,
|
||||
}
|
||||
|
||||
impl SchedulerStatsRecorder {
|
||||
/// Resolve the fixed-label metric handles for the connected engines.
|
||||
pub(crate) fn new(
|
||||
metrics: &SchedulerMetrics,
|
||||
model_name: &str,
|
||||
engines: &[ConnectedEngine],
|
||||
) -> Self {
|
||||
let engines = engines
|
||||
.iter()
|
||||
.filter_map(|engine| {
|
||||
let engine = engine.engine_id.engine_index()?;
|
||||
Some((
|
||||
engine,
|
||||
resolve_scheduler_stats_handles(metrics, model_name, engine),
|
||||
))
|
||||
})
|
||||
.collect();
|
||||
|
||||
Self { engines }
|
||||
}
|
||||
|
||||
/// Record one scheduler-stats payload for the given engine index.
|
||||
pub(crate) fn record(&self, engine_index: u32, stats: &SchedulerStats) {
|
||||
if let Some(handles) = self.engines.get(&engine_index) {
|
||||
record_scheduler_stats_with_handles(handles, stats);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve all fixed-label scheduler metrics for one engine.
|
||||
fn resolve_scheduler_stats_handles(
|
||||
metrics: &SchedulerMetrics,
|
||||
model_name: &str,
|
||||
engine: u32,
|
||||
) -> SchedulerStatsHandles {
|
||||
let labels = EngineLabels {
|
||||
model_name: model_name.to_string(),
|
||||
engine,
|
||||
};
|
||||
let capacity = WaitingReasonLabels {
|
||||
model_name: model_name.to_string(),
|
||||
engine,
|
||||
reason: WAITING_REASON_CAPACITY,
|
||||
};
|
||||
let deferred = WaitingReasonLabels {
|
||||
model_name: model_name.to_string(),
|
||||
engine,
|
||||
reason: WAITING_REASON_DEFERRED,
|
||||
};
|
||||
|
||||
SchedulerStatsHandles {
|
||||
scheduler_running: metrics.scheduler_running.get_or_create_owned(&labels),
|
||||
scheduler_waiting: metrics.scheduler_waiting.get_or_create_owned(&labels),
|
||||
scheduler_waiting_capacity: metrics
|
||||
.scheduler_waiting_by_reason
|
||||
.get_or_create_owned(&capacity),
|
||||
scheduler_waiting_deferred: metrics
|
||||
.scheduler_waiting_by_reason
|
||||
.get_or_create_owned(&deferred),
|
||||
kv_cache_usage: metrics.kv_cache_usage.get_or_create_owned(&labels),
|
||||
prefix_cache_queries: metrics.prefix_cache_queries.get_or_create_owned(&labels),
|
||||
prefix_cache_hits: metrics.prefix_cache_hits.get_or_create_owned(&labels),
|
||||
external_prefix_cache_queries: metrics
|
||||
.external_prefix_cache_queries
|
||||
.get_or_create_owned(&labels),
|
||||
external_prefix_cache_hits: metrics.external_prefix_cache_hits.get_or_create_owned(&labels),
|
||||
spec_decode_num_drafts: metrics.spec_decode_num_drafts.get_or_create_owned(&labels),
|
||||
spec_decode_num_draft_tokens: metrics
|
||||
.spec_decode_num_draft_tokens
|
||||
.get_or_create_owned(&labels),
|
||||
spec_decode_num_accepted_tokens: metrics
|
||||
.spec_decode_num_accepted_tokens
|
||||
.get_or_create_owned(&labels),
|
||||
spec_decode_num_accepted_tokens_per_pos: metrics
|
||||
.spec_decode_num_accepted_tokens_per_pos
|
||||
.clone(),
|
||||
log_stats: metrics.log_stats.get_or_create_owned(&labels),
|
||||
estimated_flops_per_gpu: metrics.estimated_flops_per_gpu.get_or_create_owned(&labels),
|
||||
estimated_read_bytes_per_gpu: metrics
|
||||
.estimated_read_bytes_per_gpu
|
||||
.get_or_create_owned(&labels),
|
||||
estimated_write_bytes_per_gpu: metrics
|
||||
.estimated_write_bytes_per_gpu
|
||||
.get_or_create_owned(&labels),
|
||||
kv_block_lifetime_seconds: metrics.kv_block_lifetime_seconds.get_or_create_owned(&labels),
|
||||
kv_block_idle_before_evict_seconds: metrics
|
||||
.kv_block_idle_before_evict_seconds
|
||||
.get_or_create_owned(&labels),
|
||||
kv_block_reuse_gap_seconds: metrics.kv_block_reuse_gap_seconds.get_or_create_owned(&labels),
|
||||
labels,
|
||||
}
|
||||
}
|
||||
|
||||
/// Record scheduler-stats values through pre-resolved metric handles.
|
||||
fn record_scheduler_stats_with_handles(handles: &SchedulerStatsHandles, stats: &SchedulerStats) {
|
||||
// Scheduler state gauges.
|
||||
handles.scheduler_running.set(stats.num_running_reqs);
|
||||
handles
|
||||
.scheduler_waiting
|
||||
.set(stats.num_waiting_reqs + stats.num_skipped_waiting_reqs);
|
||||
handles.scheduler_waiting_capacity.set(stats.num_waiting_reqs);
|
||||
handles.scheduler_waiting_deferred.set(stats.num_skipped_waiting_reqs);
|
||||
handles.kv_cache_usage.set(stats.kv_cache_usage);
|
||||
|
||||
// Prefix-cache counters, including the connector-backed external cache path.
|
||||
handles.prefix_cache_queries.inc_by(stats.prefix_cache_stats.base.queries);
|
||||
handles.prefix_cache_hits.inc_by(stats.prefix_cache_stats.base.hits);
|
||||
|
||||
if let Some(connector_prefix_cache_stats) = &stats.connector_prefix_cache_stats {
|
||||
metrics
|
||||
handles
|
||||
.external_prefix_cache_queries
|
||||
.get_or_create(&labels)
|
||||
.inc_by(connector_prefix_cache_stats.base.queries);
|
||||
metrics
|
||||
handles
|
||||
.external_prefix_cache_hits
|
||||
.get_or_create(&labels)
|
||||
.inc_by(connector_prefix_cache_stats.base.hits);
|
||||
}
|
||||
|
||||
// Speculative decoding counters.
|
||||
if let Some(spec_decoding_stats) = &stats.spec_decoding_stats {
|
||||
metrics
|
||||
.spec_decode_num_drafts
|
||||
.get_or_create(&labels)
|
||||
.inc_by(spec_decoding_stats.num_drafts);
|
||||
metrics
|
||||
handles.spec_decode_num_drafts.inc_by(spec_decoding_stats.num_drafts);
|
||||
handles
|
||||
.spec_decode_num_draft_tokens
|
||||
.get_or_create(&labels)
|
||||
.inc_by(spec_decoding_stats.num_draft_tokens);
|
||||
metrics
|
||||
handles
|
||||
.spec_decode_num_accepted_tokens
|
||||
.get_or_create(&labels)
|
||||
.inc_by(spec_decoding_stats.num_accepted_tokens);
|
||||
metrics.log_stats.get_or_create(&labels).observe_spec_decode(
|
||||
handles.log_stats.observe_spec_decode(
|
||||
spec_decoding_stats.num_drafts,
|
||||
&spec_decoding_stats.num_accepted_tokens_per_pos,
|
||||
);
|
||||
@@ -92,11 +192,11 @@ pub(crate) fn record_scheduler_stats(
|
||||
for (position, accepted_tokens) in
|
||||
spec_decoding_stats.num_accepted_tokens_per_pos.iter().copied().enumerate()
|
||||
{
|
||||
metrics
|
||||
handles
|
||||
.spec_decode_num_accepted_tokens_per_pos
|
||||
.get_or_create(&EnginePositionLabels {
|
||||
model_name: model_name.clone(),
|
||||
engine,
|
||||
model_name: handles.labels.model_name.clone(),
|
||||
engine: handles.labels.engine,
|
||||
position: position as u32,
|
||||
})
|
||||
.inc_by(accepted_tokens);
|
||||
@@ -109,22 +209,13 @@ pub(crate) fn record_scheduler_stats(
|
||||
|| perf_stats.num_read_bytes_per_gpu != 0
|
||||
|| perf_stats.num_write_bytes_per_gpu != 0)
|
||||
{
|
||||
metrics
|
||||
.estimated_flops_per_gpu
|
||||
.get_or_create(&labels)
|
||||
.inc_by(perf_stats.num_flops_per_gpu);
|
||||
metrics
|
||||
.estimated_read_bytes_per_gpu
|
||||
.get_or_create(&labels)
|
||||
.inc_by(perf_stats.num_read_bytes_per_gpu);
|
||||
metrics
|
||||
.estimated_write_bytes_per_gpu
|
||||
.get_or_create(&labels)
|
||||
.inc_by(perf_stats.num_write_bytes_per_gpu);
|
||||
handles.estimated_flops_per_gpu.inc_by(perf_stats.num_flops_per_gpu);
|
||||
handles.estimated_read_bytes_per_gpu.inc_by(perf_stats.num_read_bytes_per_gpu);
|
||||
handles.estimated_write_bytes_per_gpu.inc_by(perf_stats.num_write_bytes_per_gpu);
|
||||
}
|
||||
|
||||
if let Some(cudagraph_stats) = &stats.cudagraph_stats {
|
||||
metrics.log_stats.get_or_create(&labels).observe_cudagraph(
|
||||
handles.log_stats.observe_cudagraph(
|
||||
cudagraph_stats.num_unpadded_tokens,
|
||||
cudagraph_stats.num_padded_tokens,
|
||||
cudagraph_stats.num_paddings,
|
||||
@@ -134,16 +225,11 @@ pub(crate) fn record_scheduler_stats(
|
||||
|
||||
// Sampled KV-cache residency histograms.
|
||||
if !stats.kv_cache_eviction_events.is_empty() {
|
||||
let kv_block_lifetime_seconds = metrics.kv_block_lifetime_seconds.get_or_create(&labels);
|
||||
let kv_block_idle_before_evict_seconds =
|
||||
metrics.kv_block_idle_before_evict_seconds.get_or_create(&labels);
|
||||
let kv_block_reuse_gap_seconds = metrics.kv_block_reuse_gap_seconds.get_or_create(&labels);
|
||||
|
||||
for event in &stats.kv_cache_eviction_events {
|
||||
kv_block_lifetime_seconds.observe(event.lifetime_seconds);
|
||||
kv_block_idle_before_evict_seconds.observe(event.idle_seconds);
|
||||
handles.kv_block_lifetime_seconds.observe(event.lifetime_seconds);
|
||||
handles.kv_block_idle_before_evict_seconds.observe(event.idle_seconds);
|
||||
for reuse_gap_seconds in &event.reuse_gaps_seconds {
|
||||
kv_block_reuse_gap_seconds.observe(*reuse_gap_seconds);
|
||||
handles.kv_block_reuse_gap_seconds.observe(*reuse_gap_seconds);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -11,6 +11,21 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple};
|
||||
/// <https://github.com/vllm-project/vllm/blob/5a0a8fc1ea7542394ff315138bd5677b7b53bca1/vllm/v1/serial_utils.py#L41-L43>
|
||||
const CUSTOM_TYPE_RAW_VIEW: i8 = 3;
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename = "_ExtStruct")]
|
||||
struct MsgpackExtRef<'a>((i8, ByteSlice<'a>));
|
||||
|
||||
struct ByteSlice<'a>(&'a [u8]);
|
||||
|
||||
impl Serialize for ByteSlice<'_> {
|
||||
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
|
||||
where
|
||||
S: Serializer,
|
||||
{
|
||||
serializer.serialize_bytes(self.0)
|
||||
}
|
||||
}
|
||||
|
||||
#[easy_ext::ext(ShapeExt)]
|
||||
impl [usize] {
|
||||
/// Returned the total number of elements implied by this shape, or `None`
|
||||
@@ -184,7 +199,7 @@ impl Serialize for WireArrayData {
|
||||
match self {
|
||||
Self::AuxIndex(index) => serializer.serialize_u64(*index as u64),
|
||||
Self::RawView(bytes) => {
|
||||
Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()).serialize(serializer)
|
||||
MsgpackExtRef((CUSTOM_TYPE_RAW_VIEW, ByteSlice(bytes))).serialize(serializer)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -194,6 +209,21 @@ impl Serialize for WireArrayData {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn raw_view_serializes_as_msgpack_ext() {
|
||||
let bytes = vec![1, 2, 3, 4];
|
||||
let encoded =
|
||||
rmp_serde::to_vec_named(&WireArrayData::RawView(bytes.clone())).expect("encode");
|
||||
let expected = rmp_serde::to_vec_named(&Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes.clone()))
|
||||
.expect("encode expected");
|
||||
|
||||
assert_eq!(encoded, expected);
|
||||
assert_eq!(
|
||||
rmpv::decode::read_value(&mut std::io::Cursor::new(encoded)).expect("decode"),
|
||||
Value::Ext(CUSTOM_TYPE_RAW_VIEW, bytes)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn constructors_build_raw_view_tensors() {
|
||||
let f32_tensor = WireNdArray::from_f32(vec![2], vec![1.0, 2.5]).unwrap();
|
||||
|
||||
+12
-4
@@ -14,6 +14,7 @@ pub use output::{
|
||||
GenerateOutputStreamExt, GeneratePromptInfo, TokenUsage,
|
||||
};
|
||||
pub use request::GenerateRequest;
|
||||
pub use request_metrics::current_unix_timestamp_secs;
|
||||
pub use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob};
|
||||
|
||||
use crate::inflight::InflightRequests;
|
||||
@@ -88,14 +89,21 @@ impl Llm {
|
||||
// Record internal engine-core request ID in the current tracing span.
|
||||
Span::current().record("engine_request_id", &internal_request_id);
|
||||
|
||||
let arrival_time = prepared.engine_request.arrival_time;
|
||||
let max_tokens_param =
|
||||
(prepared.engine_request.sampling_params.as_ref()).map(|p| p.max_tokens);
|
||||
let prompt_len = prepared.prompt_token_ids().len() as u32;
|
||||
|
||||
let stream = self.client.call(prepared.engine_request).await?;
|
||||
|
||||
let request_metrics = RequestMetricsTracker::new(
|
||||
self.client.model_name().to_string(),
|
||||
prepared.engine_request.arrival_time,
|
||||
prepared.prompt_token_ids().len() as u32,
|
||||
(prepared.engine_request.sampling_params.as_ref()).map(|p| p.max_tokens),
|
||||
stream.engine_index(),
|
||||
arrival_time,
|
||||
prompt_len,
|
||||
max_tokens_param,
|
||||
1,
|
||||
);
|
||||
let stream = self.client.call(prepared.engine_request).await?;
|
||||
let guard = self.inflight.track(external_request_id, internal_request_id);
|
||||
|
||||
Ok(GenerateOutputStream::new(
|
||||
|
||||
@@ -248,12 +248,7 @@ impl Stream for GenerateOutputStream {
|
||||
};
|
||||
|
||||
let received_at = current_unix_timestamp_secs();
|
||||
self.request_metrics.observe_output(
|
||||
raw.engine_index,
|
||||
raw.timestamp,
|
||||
received_at,
|
||||
&raw.output,
|
||||
);
|
||||
self.request_metrics.observe_output(raw.timestamp, received_at, &raw.output);
|
||||
|
||||
let raw = raw.output;
|
||||
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use uuid::Uuid;
|
||||
use vllm_engine_core_client::protocol::lora::LoraRequest;
|
||||
@@ -8,6 +7,7 @@ use vllm_engine_core_client::protocol::request::{EngineCoreRequest, ReasoningPar
|
||||
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
use crate::request_metrics::current_unix_timestamp_secs;
|
||||
|
||||
/// Tokenized decoder-only generate request accepted by [`crate::Llm`].
|
||||
///
|
||||
@@ -30,8 +30,9 @@ pub struct GenerateRequest {
|
||||
pub mm_features: Option<MmFeatures>,
|
||||
/// Unix timestamp, in seconds, when this request arrived at the frontend.
|
||||
///
|
||||
/// When omitted, the Rust frontend fills it immediately before sending the
|
||||
/// request to engine-core, matching Python's default arrival-time behavior.
|
||||
/// Stamped at the frontend entry, before render and tokenization, to match
|
||||
/// Python's renderer-entry arrival_time. When omitted, it is filled as a
|
||||
/// fallback before the request is sent to engine-core.
|
||||
pub arrival_time: Option<f64>,
|
||||
/// Optional salt used to partition prefix-cache entries for this request.
|
||||
pub cache_salt: Option<String>,
|
||||
@@ -122,13 +123,6 @@ impl PreparedGenerateRequest {
|
||||
}
|
||||
}
|
||||
|
||||
fn current_unix_timestamp_secs() -> f64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("system clock is before unix epoch")
|
||||
.as_secs_f64()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
+151
-143
@@ -5,15 +5,12 @@ use vllm_engine_core_client::protocol::output::{
|
||||
};
|
||||
use vllm_engine_core_client::protocol::stats::PrefillStats;
|
||||
use vllm_metrics::{
|
||||
EngineLabels, FinishedReasonLabels, METRICS, PromptTokenSourceLabels, RequestMetrics,
|
||||
EngineLabels, Family, FinishedReasonLabels, HistogramMetric, METRICS, PromptTokenSourceLabels,
|
||||
U64Counter,
|
||||
};
|
||||
|
||||
use crate::FinishReason;
|
||||
|
||||
fn metrics() -> &'static RequestMetrics {
|
||||
&METRICS.request
|
||||
}
|
||||
|
||||
const PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE: &str = "local_compute";
|
||||
const PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT: &str = "local_cache_hit";
|
||||
const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer";
|
||||
@@ -29,9 +26,11 @@ const PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER: &str = "external_kv_transfer";
|
||||
///
|
||||
/// Original Python update flow:
|
||||
/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/engine/output_processor.py#L600-L677>
|
||||
#[derive(Debug, Clone)]
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct RequestMetricsTracker {
|
||||
model_name: String,
|
||||
/// Cached request metric handles for this request's model and engine index.
|
||||
handles: RequestMetricHandles,
|
||||
|
||||
arrival_time: f64,
|
||||
prompt_len: u32,
|
||||
max_tokens_param: Option<u32>,
|
||||
@@ -44,7 +43,38 @@ pub(crate) struct RequestMetricsTracker {
|
||||
first_token_latency: f64,
|
||||
num_generation_tokens: u32,
|
||||
latest_num_cached_tokens: u32,
|
||||
last_seen_engine_index: u32,
|
||||
}
|
||||
|
||||
/// Cached request metric handles for one model and engine index.
|
||||
#[derive(Clone)]
|
||||
struct RequestMetricHandles {
|
||||
labels: EngineLabels,
|
||||
|
||||
// Request-derived counters.
|
||||
num_preemptions: U64Counter,
|
||||
prompt_tokens: U64Counter,
|
||||
prompt_tokens_local_compute: U64Counter,
|
||||
prompt_tokens_local_cache_hit: U64Counter,
|
||||
prompt_tokens_external_kv_transfer: U64Counter,
|
||||
prompt_tokens_cached: U64Counter,
|
||||
generation_tokens: U64Counter,
|
||||
|
||||
// Request lifecycle counters and histograms.
|
||||
request_success: Family<FinishedReasonLabels, U64Counter>,
|
||||
request_prompt_tokens: HistogramMetric,
|
||||
request_generation_tokens: HistogramMetric,
|
||||
request_max_num_generation_tokens: HistogramMetric,
|
||||
request_params_max_tokens: HistogramMetric,
|
||||
request_params_n: HistogramMetric,
|
||||
request_prefill_kv_computed_tokens: HistogramMetric,
|
||||
time_to_first_token_seconds: HistogramMetric,
|
||||
inter_token_latency_seconds: HistogramMetric,
|
||||
e2e_request_latency_seconds: HistogramMetric,
|
||||
request_queue_time_seconds: HistogramMetric,
|
||||
request_prefill_time_seconds: HistogramMetric,
|
||||
request_decode_time_seconds: HistogramMetric,
|
||||
request_inference_time_seconds: HistogramMetric,
|
||||
request_time_per_output_token_seconds: HistogramMetric,
|
||||
}
|
||||
|
||||
impl RequestMetricsTracker {
|
||||
@@ -52,13 +82,14 @@ impl RequestMetricsTracker {
|
||||
/// context.
|
||||
pub(crate) fn new(
|
||||
model_name: String,
|
||||
engine_index: u32,
|
||||
arrival_time: f64,
|
||||
prompt_len: u32,
|
||||
max_tokens_param: Option<u32>,
|
||||
n_param: u32,
|
||||
) -> Self {
|
||||
Self {
|
||||
model_name,
|
||||
handles: resolve_request_metric_handles(&model_name, engine_index),
|
||||
arrival_time,
|
||||
prompt_len,
|
||||
max_tokens_param,
|
||||
@@ -71,7 +102,6 @@ impl RequestMetricsTracker {
|
||||
first_token_latency: 0.0,
|
||||
num_generation_tokens: 0,
|
||||
latest_num_cached_tokens: 0,
|
||||
last_seen_engine_index: 0,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -81,23 +111,18 @@ impl RequestMetricsTracker {
|
||||
/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L331-L384>
|
||||
pub(crate) fn observe_output(
|
||||
&mut self,
|
||||
engine_index: u32,
|
||||
batch_timestamp: f64,
|
||||
received_at: f64,
|
||||
output: &EngineCoreOutput,
|
||||
) {
|
||||
self.last_seen_engine_index = engine_index;
|
||||
if let Some(prefill_stats) = &output.prefill_stats {
|
||||
self.latest_num_cached_tokens = prefill_stats.num_cached_tokens;
|
||||
}
|
||||
self.num_generation_tokens += output.new_token_ids.len() as u32;
|
||||
metrics()
|
||||
.generation_tokens
|
||||
.get_or_create(&engine_labels(&self.model_name, engine_index))
|
||||
.inc_by(output.new_token_ids.len() as u64);
|
||||
self.handles.generation_tokens.inc_by(output.new_token_ids.len() as u64);
|
||||
|
||||
if let Some(events) = &output.events {
|
||||
self.observe_events(engine_index, events);
|
||||
self.observe_events(events);
|
||||
}
|
||||
|
||||
// Only outputs that actually carry tokens drive token-timing metrics.
|
||||
@@ -107,22 +132,16 @@ impl RequestMetricsTracker {
|
||||
if !output.new_token_ids.is_empty() {
|
||||
if self.is_prefilling {
|
||||
if let Some(prefill_stats) = &output.prefill_stats {
|
||||
record_prompt_tokens(&self.model_name, engine_index, prefill_stats);
|
||||
self.record_prompt_tokens(prefill_stats);
|
||||
}
|
||||
self.first_token_latency = received_at - self.arrival_time;
|
||||
observe_time_to_first_token_seconds(
|
||||
&self.model_name,
|
||||
engine_index,
|
||||
self.first_token_latency,
|
||||
);
|
||||
self.handles.time_to_first_token_seconds.observe(self.first_token_latency);
|
||||
self.first_token_ts = batch_timestamp;
|
||||
self.is_prefilling = false;
|
||||
} else if self.last_token_ts > 0.0 {
|
||||
observe_inter_token_latency_seconds(
|
||||
&self.model_name,
|
||||
engine_index,
|
||||
batch_timestamp - self.last_token_ts,
|
||||
);
|
||||
self.handles
|
||||
.inter_token_latency_seconds
|
||||
.observe(batch_timestamp - self.last_token_ts);
|
||||
}
|
||||
|
||||
self.last_token_ts = batch_timestamp;
|
||||
@@ -135,7 +154,6 @@ impl RequestMetricsTracker {
|
||||
/// Original Python finished-request stats:
|
||||
/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L222-L237>
|
||||
pub(crate) fn record_finished(&self, received_at: f64, finish_reason: FinishReason) {
|
||||
let labels = engine_labels(&self.model_name, self.last_seen_engine_index);
|
||||
let prefill_kv_computed_tokens =
|
||||
self.prompt_len.saturating_sub(self.latest_num_cached_tokens);
|
||||
let e2e_latency_seconds = received_at - self.arrival_time;
|
||||
@@ -150,57 +168,47 @@ impl RequestMetricsTracker {
|
||||
0.0
|
||||
};
|
||||
|
||||
record_request_success(&self.model_name, self.last_seen_engine_index, finish_reason);
|
||||
metrics()
|
||||
.request_prompt_tokens
|
||||
.get_or_create(&labels)
|
||||
.observe(self.prompt_len as f64);
|
||||
metrics()
|
||||
self.record_request_success(finish_reason);
|
||||
|
||||
self.handles.request_prompt_tokens.observe(self.prompt_len as f64);
|
||||
self.handles
|
||||
.request_generation_tokens
|
||||
.get_or_create(&labels)
|
||||
.observe(self.num_generation_tokens as f64);
|
||||
metrics()
|
||||
self.handles
|
||||
.request_max_num_generation_tokens
|
||||
.get_or_create(&labels)
|
||||
.observe(self.num_generation_tokens as f64);
|
||||
if let Some(max_tokens_param) = self.max_tokens_param {
|
||||
metrics()
|
||||
.request_params_max_tokens
|
||||
.get_or_create(&labels)
|
||||
.observe(max_tokens_param as f64);
|
||||
self.handles.request_params_max_tokens.observe(max_tokens_param as f64);
|
||||
}
|
||||
metrics().request_params_n.get_or_create(&labels).observe(self.n_param as f64);
|
||||
metrics()
|
||||
self.handles.request_params_n.observe(self.n_param as f64);
|
||||
self.handles
|
||||
.request_prefill_kv_computed_tokens
|
||||
.get_or_create(&labels)
|
||||
.observe(prefill_kv_computed_tokens as f64);
|
||||
metrics()
|
||||
.e2e_request_latency_seconds
|
||||
.get_or_create(&labels)
|
||||
.observe(e2e_latency_seconds);
|
||||
metrics()
|
||||
.request_queue_time_seconds
|
||||
.get_or_create(&labels)
|
||||
.observe(queue_time_seconds);
|
||||
metrics()
|
||||
.request_prefill_time_seconds
|
||||
.get_or_create(&labels)
|
||||
.observe(prefill_time_seconds);
|
||||
metrics()
|
||||
.request_decode_time_seconds
|
||||
.get_or_create(&labels)
|
||||
.observe(decode_time_seconds);
|
||||
metrics()
|
||||
.request_inference_time_seconds
|
||||
.get_or_create(&labels)
|
||||
.observe(inference_time_seconds);
|
||||
metrics()
|
||||
self.handles.e2e_request_latency_seconds.observe(e2e_latency_seconds);
|
||||
self.handles.request_queue_time_seconds.observe(queue_time_seconds);
|
||||
self.handles.request_prefill_time_seconds.observe(prefill_time_seconds);
|
||||
self.handles.request_decode_time_seconds.observe(decode_time_seconds);
|
||||
self.handles.request_inference_time_seconds.observe(inference_time_seconds);
|
||||
self.handles
|
||||
.request_time_per_output_token_seconds
|
||||
.get_or_create(&labels)
|
||||
.observe(time_per_output_token_seconds);
|
||||
}
|
||||
|
||||
fn observe_events(&mut self, engine_index: u32, events: &[EngineCoreEvent]) {
|
||||
/// Record prompt token counters through cached metric handles.
|
||||
fn record_prompt_tokens(&self, prefill_stats: &PrefillStats) {
|
||||
let computed = prefill_stats.num_computed_tokens as u64;
|
||||
let local_cache_hit = prefill_stats.num_local_cached_tokens as u64;
|
||||
let external_kv_transfer = prefill_stats.num_external_cached_tokens as u64;
|
||||
|
||||
self.handles.prompt_tokens.inc_by(prefill_stats.num_prompt_tokens as u64);
|
||||
self.handles.prompt_tokens_local_compute.inc_by(computed);
|
||||
self.handles.prompt_tokens_local_cache_hit.inc_by(local_cache_hit);
|
||||
self.handles.prompt_tokens_external_kv_transfer.inc_by(external_kv_transfer);
|
||||
self.handles.prompt_tokens_cached.inc_by(prefill_stats.num_cached_tokens as u64);
|
||||
}
|
||||
|
||||
/// Record request event counters through cached metric handles.
|
||||
fn observe_events(&mut self, events: &[EngineCoreEvent]) {
|
||||
for event in events {
|
||||
match event.r#type {
|
||||
EngineCoreEventType::Queued => {
|
||||
@@ -212,46 +220,86 @@ impl RequestMetricsTracker {
|
||||
}
|
||||
}
|
||||
EngineCoreEventType::Preempted => {
|
||||
metrics()
|
||||
.num_preemptions
|
||||
.get_or_create(&engine_labels(&self.model_name, engine_index))
|
||||
.inc();
|
||||
self.handles.num_preemptions.inc();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn engine_labels(model_name: &str, engine: u32) -> EngineLabels {
|
||||
EngineLabels {
|
||||
model_name: model_name.to_string(),
|
||||
engine,
|
||||
/// Increment the request-success counter for the terminal finish reason.
|
||||
fn record_request_success(&self, finish_reason: FinishReason) {
|
||||
self.handles
|
||||
.request_success
|
||||
.get_or_create(&FinishedReasonLabels {
|
||||
model_name: self.handles.labels.model_name.clone(),
|
||||
engine: self.handles.labels.engine,
|
||||
finished_reason: finish_reason.as_str(),
|
||||
})
|
||||
.inc();
|
||||
}
|
||||
}
|
||||
|
||||
fn observe_time_to_first_token_seconds(model_name: &str, engine: u32, seconds: f64) {
|
||||
metrics()
|
||||
.time_to_first_token_seconds
|
||||
.get_or_create(&engine_labels(model_name, engine))
|
||||
.observe(seconds);
|
||||
}
|
||||
/// Resolve fixed request metric handles for one model and engine index.
|
||||
fn resolve_request_metric_handles(model_name: &str, engine: u32) -> RequestMetricHandles {
|
||||
let metrics = &METRICS.request;
|
||||
let labels = EngineLabels {
|
||||
model_name: model_name.to_string(),
|
||||
engine,
|
||||
};
|
||||
|
||||
fn observe_inter_token_latency_seconds(model_name: &str, engine: u32, seconds: f64) {
|
||||
metrics()
|
||||
.inter_token_latency_seconds
|
||||
.get_or_create(&engine_labels(model_name, engine))
|
||||
.observe(seconds);
|
||||
}
|
||||
|
||||
fn record_request_success(model_name: &str, engine: u32, finish_reason: FinishReason) {
|
||||
metrics()
|
||||
.request_success
|
||||
.get_or_create(&FinishedReasonLabels {
|
||||
model_name: model_name.to_string(),
|
||||
engine,
|
||||
finished_reason: finish_reason.as_str(),
|
||||
})
|
||||
.inc();
|
||||
RequestMetricHandles {
|
||||
num_preemptions: metrics.num_preemptions.get_or_create_owned(&labels),
|
||||
prompt_tokens: metrics.prompt_tokens.get_or_create_owned(&labels),
|
||||
prompt_tokens_local_compute: metrics.prompt_tokens_by_source.get_or_create_owned(
|
||||
&prompt_token_source_labels(model_name, engine, PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE),
|
||||
),
|
||||
prompt_tokens_local_cache_hit: metrics.prompt_tokens_by_source.get_or_create_owned(
|
||||
&prompt_token_source_labels(model_name, engine, PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT),
|
||||
),
|
||||
prompt_tokens_external_kv_transfer: metrics.prompt_tokens_by_source.get_or_create_owned(
|
||||
&prompt_token_source_labels(
|
||||
model_name,
|
||||
engine,
|
||||
PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER,
|
||||
),
|
||||
),
|
||||
prompt_tokens_cached: metrics.prompt_tokens_cached.get_or_create_owned(&labels),
|
||||
generation_tokens: metrics.generation_tokens.get_or_create_owned(&labels),
|
||||
request_success: metrics.request_success.clone(),
|
||||
request_prompt_tokens: metrics.request_prompt_tokens.get_or_create_owned(&labels),
|
||||
request_generation_tokens: metrics.request_generation_tokens.get_or_create_owned(&labels),
|
||||
request_max_num_generation_tokens: metrics
|
||||
.request_max_num_generation_tokens
|
||||
.get_or_create_owned(&labels),
|
||||
request_params_max_tokens: metrics.request_params_max_tokens.get_or_create_owned(&labels),
|
||||
request_params_n: metrics.request_params_n.get_or_create_owned(&labels),
|
||||
request_prefill_kv_computed_tokens: metrics
|
||||
.request_prefill_kv_computed_tokens
|
||||
.get_or_create_owned(&labels),
|
||||
time_to_first_token_seconds: metrics
|
||||
.time_to_first_token_seconds
|
||||
.get_or_create_owned(&labels),
|
||||
inter_token_latency_seconds: metrics
|
||||
.inter_token_latency_seconds
|
||||
.get_or_create_owned(&labels),
|
||||
e2e_request_latency_seconds: metrics
|
||||
.e2e_request_latency_seconds
|
||||
.get_or_create_owned(&labels),
|
||||
request_queue_time_seconds: metrics.request_queue_time_seconds.get_or_create_owned(&labels),
|
||||
request_prefill_time_seconds: metrics
|
||||
.request_prefill_time_seconds
|
||||
.get_or_create_owned(&labels),
|
||||
request_decode_time_seconds: metrics
|
||||
.request_decode_time_seconds
|
||||
.get_or_create_owned(&labels),
|
||||
request_inference_time_seconds: metrics
|
||||
.request_inference_time_seconds
|
||||
.get_or_create_owned(&labels),
|
||||
request_time_per_output_token_seconds: metrics
|
||||
.request_time_per_output_token_seconds
|
||||
.get_or_create_owned(&labels),
|
||||
labels,
|
||||
}
|
||||
}
|
||||
|
||||
fn prompt_token_source_labels(
|
||||
@@ -266,45 +314,6 @@ fn prompt_token_source_labels(
|
||||
}
|
||||
}
|
||||
|
||||
fn record_prompt_tokens(model_name: &str, engine: u32, prefill_stats: &PrefillStats) {
|
||||
let computed = prefill_stats.num_computed_tokens as u64;
|
||||
let local_cache_hit = prefill_stats.num_local_cached_tokens as u64;
|
||||
let external_kv_transfer = prefill_stats.num_external_cached_tokens as u64;
|
||||
|
||||
metrics()
|
||||
.prompt_tokens
|
||||
.get_or_create(&engine_labels(model_name, engine))
|
||||
.inc_by(prefill_stats.num_prompt_tokens as u64);
|
||||
metrics()
|
||||
.prompt_tokens_by_source
|
||||
.get_or_create(&prompt_token_source_labels(
|
||||
model_name,
|
||||
engine,
|
||||
PROMPT_TOKEN_SOURCE_LOCAL_COMPUTE,
|
||||
))
|
||||
.inc_by(computed);
|
||||
metrics()
|
||||
.prompt_tokens_by_source
|
||||
.get_or_create(&prompt_token_source_labels(
|
||||
model_name,
|
||||
engine,
|
||||
PROMPT_TOKEN_SOURCE_LOCAL_CACHE_HIT,
|
||||
))
|
||||
.inc_by(local_cache_hit);
|
||||
metrics()
|
||||
.prompt_tokens_by_source
|
||||
.get_or_create(&prompt_token_source_labels(
|
||||
model_name,
|
||||
engine,
|
||||
PROMPT_TOKEN_SOURCE_EXTERNAL_KV_TRANSFER,
|
||||
))
|
||||
.inc_by(external_kv_transfer);
|
||||
metrics()
|
||||
.prompt_tokens_cached
|
||||
.get_or_create(&engine_labels(model_name, engine))
|
||||
.inc_by(prefill_stats.num_cached_tokens as u64);
|
||||
}
|
||||
|
||||
fn diff_or_zero(end: f64, start: f64) -> f64 {
|
||||
if end > 0.0 && start > 0.0 && end >= start {
|
||||
end - start
|
||||
@@ -321,7 +330,7 @@ fn diff_or_zero(end: f64, start: f64) -> f64 {
|
||||
///
|
||||
/// Original Python request timestamp source:
|
||||
/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/v1/metrics/stats.py#L206-L216>
|
||||
pub(crate) fn current_unix_timestamp_secs() -> f64 {
|
||||
pub fn current_unix_timestamp_secs() -> f64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("system clock is before unix epoch")
|
||||
@@ -337,10 +346,10 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn tracker_updates_timing_state_across_prefill_decode_and_finish() {
|
||||
let mut tracker = RequestMetricsTracker::new("model".to_string(), 100.0, 64, Some(128), 1);
|
||||
let mut tracker =
|
||||
RequestMetricsTracker::new("model".to_string(), 2, 100.0, 64, Some(128), 1);
|
||||
|
||||
tracker.observe_output(
|
||||
2,
|
||||
10.0,
|
||||
100.2,
|
||||
&vllm_engine_core_client::protocol::output::EngineCoreOutput {
|
||||
@@ -368,7 +377,6 @@ mod tests {
|
||||
},
|
||||
);
|
||||
tracker.observe_output(
|
||||
2,
|
||||
11.5,
|
||||
100.4,
|
||||
&vllm_engine_core_client::protocol::output::EngineCoreOutput {
|
||||
@@ -384,7 +392,7 @@ mod tests {
|
||||
);
|
||||
|
||||
assert!(!tracker.is_prefilling);
|
||||
assert_eq!(tracker.last_seen_engine_index, 2);
|
||||
assert_eq!(tracker.handles.labels.engine, 2);
|
||||
assert_eq!(tracker.num_generation_tokens, 3);
|
||||
assert_eq!(tracker.queued_ts, 8.0);
|
||||
assert_eq!(tracker.scheduled_ts, 9.0);
|
||||
|
||||
@@ -17,7 +17,7 @@ use vllm_engine_core_client::protocol::request::EngineCoreRequest;
|
||||
use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;
|
||||
use vllm_engine_core_client::protocol::stats::PrefillStats;
|
||||
use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task};
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig};
|
||||
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId};
|
||||
use vllm_llm::{
|
||||
Error, FinishReason, GenerateOutputStreamExt as _, GeneratePromptInfo, GenerateRequest, Llm,
|
||||
};
|
||||
@@ -699,7 +699,7 @@ async fn abort_by_external_id_aborts_all_internal_requests() {
|
||||
async fn generate_records_request_metrics_in_prometheus_output() {
|
||||
let ipc = IpcNamespace::new().unwrap();
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
let engine_id = b"engine-metrics".to_vec();
|
||||
let engine_id = EngineId::from_engine_index(4);
|
||||
let model_name = request_metrics_model_name("metrics-model");
|
||||
|
||||
let (shutdown_tx, engine_task) = spawn_mock_engine_task(
|
||||
@@ -832,7 +832,7 @@ async fn generate_records_request_metrics_in_prometheus_output() {
|
||||
async fn dropping_stream_records_abort_terminal_request_metrics() {
|
||||
let ipc = IpcNamespace::new().unwrap();
|
||||
let handshake_address = ipc.handshake_endpoint();
|
||||
let engine_id = b"engine-metrics-drop".to_vec();
|
||||
let engine_id = EngineId::from_engine_index(5);
|
||||
let model_name = request_metrics_model_name("metrics-drop-model");
|
||||
|
||||
let (shutdown_tx, engine_task) = spawn_mock_engine_task(
|
||||
|
||||
@@ -4,7 +4,7 @@ use std::sync::atomic::AtomicU64;
|
||||
|
||||
use prometheus_client::encoding::text::encode;
|
||||
use prometheus_client::metrics::counter::Counter;
|
||||
use prometheus_client::metrics::family::Family;
|
||||
pub use prometheus_client::metrics::family::Family;
|
||||
use prometheus_client::metrics::gauge::Gauge;
|
||||
use prometheus_client::metrics::histogram::Histogram;
|
||||
use prometheus_client::registry::Registry;
|
||||
@@ -23,6 +23,8 @@ pub use scheduler::*;
|
||||
pub type U64Counter = Counter<u64, AtomicU64>;
|
||||
pub type U64Gauge = Gauge<u64, AtomicU64>;
|
||||
pub type F64Gauge = Gauge<f64, AtomicU64>;
|
||||
/// Histogram metric handle cloned out of a Prometheus family.
|
||||
pub type HistogramMetric = Histogram;
|
||||
pub(crate) type HistogramFamily = Family<EngineLabels, Histogram, fn() -> Histogram>;
|
||||
|
||||
/// Shared Prometheus registry for frontend metrics.
|
||||
|
||||
@@ -94,6 +94,7 @@ pub fn to_text_request(
|
||||
data_parallel_rank: None,
|
||||
reasoning_parser_kwargs: None,
|
||||
lora_request: None,
|
||||
arrival_time: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -70,6 +70,7 @@ pub(super) fn prepare_generate_request(
|
||||
data_parallel_rank: ctx.data_parallel_rank,
|
||||
reasoning_parser_kwargs: None,
|
||||
lora_request: lora_resolution.lora_request.clone(),
|
||||
arrival_time: None,
|
||||
};
|
||||
|
||||
Ok(PreparedRequest {
|
||||
|
||||
@@ -29,7 +29,9 @@ pub struct GenerateRequest {
|
||||
impl Normalizable for GenerateRequest {}
|
||||
|
||||
/// Mirrors the Python vLLM `GenerateResponseChoice` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
///
|
||||
/// Do not skip serializing `None` fields here: non-streaming response types
|
||||
/// should serialize `None` as explicit `null`.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct GenerateResponseChoice {
|
||||
pub index: u32,
|
||||
@@ -58,7 +60,6 @@ pub(super) struct GenerateStreamResponse {
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `GenerateResponse` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct GenerateResponse {
|
||||
pub request_id: String,
|
||||
@@ -68,7 +69,6 @@ pub(super) struct GenerateResponse {
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `Logprob` class used in prompt-logprobs payloads.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct GenerateLogprob {
|
||||
pub logprob: f32,
|
||||
|
||||
@@ -211,7 +211,7 @@ async fn collect_chat_completion(
|
||||
Some(prefix) => Some(format!("{prefix}{}", message.text())),
|
||||
None => Some(message.text()).filter(|t| !t.is_empty()),
|
||||
},
|
||||
tool_calls: Some(tool_calls).filter(|calls| !calls.is_empty()),
|
||||
tool_calls,
|
||||
reasoning: if include_reasoning { reasoning } else { None },
|
||||
},
|
||||
logprobs,
|
||||
|
||||
@@ -830,7 +830,7 @@ mod tests {
|
||||
let message = ChatCompletionMessage {
|
||||
role: AssistantRole,
|
||||
content: Some("answer".to_string()),
|
||||
tool_calls: None,
|
||||
tool_calls: Vec::new(),
|
||||
reasoning: Some("inner".to_string()),
|
||||
};
|
||||
let message_json = serde_json::to_value(message).expect("message serializes");
|
||||
|
||||
@@ -331,7 +331,9 @@ impl Normalizable for ChatCompletionRequest {
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `ChatCompletionResponse` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
///
|
||||
/// Do not skip serializing `None` fields here: non-streaming response types
|
||||
/// should serialize `None` as explicit `null`.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct ChatCompletionResponse {
|
||||
pub id: String,
|
||||
@@ -347,7 +349,6 @@ pub(super) struct ChatCompletionResponse {
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `ChatCompletionResponseChoice` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct ChatCompletionChoice {
|
||||
pub index: u32,
|
||||
@@ -370,12 +371,12 @@ impl fmt::Display for AssistantRole {
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM response `ChatMessage` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct ChatCompletionMessage {
|
||||
pub role: AssistantRole,
|
||||
pub content: Option<String>,
|
||||
pub tool_calls: Option<Vec<ToolCall>>,
|
||||
#[serde(skip_serializing_if = "Vec::is_empty")]
|
||||
pub tool_calls: Vec<ToolCall>,
|
||||
pub reasoning: Option<String>,
|
||||
}
|
||||
|
||||
|
||||
@@ -144,6 +144,7 @@ pub(super) fn prepare_completion_request(
|
||||
data_parallel_rank: ctx.data_parallel_rank,
|
||||
reasoning_parser_kwargs: None,
|
||||
lora_request: lora_resolution.lora_request.clone(),
|
||||
arrival_time: None,
|
||||
};
|
||||
|
||||
Ok(PreparedRequest {
|
||||
|
||||
@@ -196,7 +196,9 @@ impl Normalizable for CompletionRequest {
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `CompletionResponse` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
///
|
||||
/// Do not skip serializing `None` fields here: non-streaming response types
|
||||
/// should serialize `None` as explicit `null`.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct CompletionResponse {
|
||||
pub id: String,
|
||||
@@ -210,7 +212,6 @@ pub(super) struct CompletionResponse {
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `CompletionResponseChoice` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub(super) struct CompletionChoice {
|
||||
pub index: u32,
|
||||
|
||||
@@ -311,7 +311,9 @@ pub enum MessageContent {
|
||||
// ============================================================================
|
||||
|
||||
/// Mirrors the Python vLLM `UsageInfo` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
///
|
||||
/// Do not skip serializing `None` fields here: non-streaming response types
|
||||
/// should serialize `None` as explicit `null`.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct Usage {
|
||||
pub prompt_tokens: usize,
|
||||
@@ -402,14 +404,12 @@ pub struct LogProbs {
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `ChatCompletionLogProbs` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ChatLogProbs {
|
||||
pub content: Option<Vec<ChatLogProbsContent>>,
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `ChatCompletionLogProbsContent` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ChatLogProbsContent {
|
||||
pub token: String,
|
||||
@@ -419,7 +419,6 @@ pub struct ChatLogProbsContent {
|
||||
}
|
||||
|
||||
/// Mirrors the Python vLLM `ChatCompletionLogProb` class.
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct TopLogProb {
|
||||
pub token: String,
|
||||
@@ -436,7 +435,6 @@ pub struct ErrorResponse {
|
||||
pub error: ErrorDetail,
|
||||
}
|
||||
|
||||
#[serde_with::skip_serializing_none]
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub struct ErrorDetail {
|
||||
pub message: String,
|
||||
|
||||
@@ -2191,6 +2191,28 @@ async fn non_stream_chat_returns_json_response() {
|
||||
assert_eq!(json["usage"]["prompt_tokens"], 22);
|
||||
assert_eq!(json["usage"]["completion_tokens"], 3);
|
||||
assert_eq!(json["usage"]["total_tokens"], 25);
|
||||
|
||||
// Unset optional fields are serialized as explicit `null` on
|
||||
// non-streaming responses...
|
||||
let response_object = json.as_object().expect("response object");
|
||||
let choice = json["choices"][0].as_object().expect("choice object");
|
||||
let message = choice["message"].as_object().expect("message object");
|
||||
for (object, key) in [
|
||||
(response_object, "system_fingerprint"),
|
||||
(response_object, "prompt_token_ids"),
|
||||
(response_object, "kv_transfer_params"),
|
||||
(choice, "logprobs"),
|
||||
(choice, "stop_reason"),
|
||||
(choice, "token_ids"),
|
||||
(message, "reasoning"),
|
||||
] {
|
||||
assert!(
|
||||
object.contains_key(key) && object[key].is_null(),
|
||||
"expected explicit null `{key}`: {json}"
|
||||
);
|
||||
}
|
||||
// ...except `tool_calls`, which Python pops from the payload when empty.
|
||||
assert!(!message.contains_key("tool_calls"), "{json}");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
@@ -2956,6 +2978,25 @@ async fn non_stream_completions_return_json_response() {
|
||||
assert_eq!(json["choices"][0]["text"], "hi");
|
||||
assert_eq!(json["choices"][0]["finish_reason"], "stop");
|
||||
assert_eq!(json["usage"]["completion_tokens"], 3);
|
||||
|
||||
// Unset optional fields are serialized as explicit `null` on
|
||||
// non-streaming responses.
|
||||
let response_object = json.as_object().expect("response object");
|
||||
let choice = json["choices"][0].as_object().expect("choice object");
|
||||
for (object, key) in [
|
||||
(response_object, "system_fingerprint"),
|
||||
(response_object, "kv_transfer_params"),
|
||||
(choice, "logprobs"),
|
||||
(choice, "stop_reason"),
|
||||
(choice, "prompt_logprobs"),
|
||||
(choice, "token_ids"),
|
||||
(choice, "prompt_token_ids"),
|
||||
] {
|
||||
assert!(
|
||||
object.contains_key(key) && object[key].is_null(),
|
||||
"expected explicit null `{key}`: {json}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
@@ -4371,10 +4412,10 @@ async fn include_reasoning_false_suppresses_reasoning_in_non_stream_chat() {
|
||||
let json: serde_json::Value = serde_json::from_str(&text).expect("decode json");
|
||||
|
||||
assert_eq!(json["choices"][0]["message"]["content"], "answer");
|
||||
// Suppressed fields are serialized as explicit `null` on non-streaming
|
||||
// responses.
|
||||
assert!(
|
||||
json["choices"][0]["message"]
|
||||
.as_object()
|
||||
.is_some_and(|message| !message.contains_key("reasoning")),
|
||||
json["choices"][0]["message"]["reasoning"].is_null(),
|
||||
"{text}"
|
||||
);
|
||||
}
|
||||
@@ -4476,14 +4517,14 @@ async fn include_reasoning_false_suppresses_non_stream_output_metadata() {
|
||||
let choice = json["choices"][0].as_object().expect("choice object");
|
||||
|
||||
assert_eq!(json["choices"][0]["message"]["content"], "answer");
|
||||
// Suppressed fields are serialized as explicit `null` on non-streaming
|
||||
// responses.
|
||||
assert!(
|
||||
json["choices"][0]["message"]
|
||||
.as_object()
|
||||
.is_some_and(|message| !message.contains_key("reasoning")),
|
||||
json["choices"][0]["message"]["reasoning"].is_null(),
|
||||
"{text}"
|
||||
);
|
||||
assert!(!choice.contains_key("logprobs"), "{text}");
|
||||
assert!(!choice.contains_key("token_ids"), "{text}");
|
||||
assert!(choice["logprobs"].is_null(), "{text}");
|
||||
assert!(choice["token_ids"].is_null(), "{text}");
|
||||
assert!(json["prompt_token_ids"].is_array(), "{text}");
|
||||
}
|
||||
|
||||
|
||||
@@ -102,12 +102,13 @@ pub struct DetokenizeRequest {
|
||||
pub tokens: Vec<u32>,
|
||||
}
|
||||
|
||||
/// Do not skip serializing `None` fields here: non-streaming response types
|
||||
/// should serialize `None` as explicit `null`.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct TokenizeResponse {
|
||||
pub count: usize,
|
||||
pub max_model_len: u32,
|
||||
pub tokens: Vec<u32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub token_strs: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
|
||||
@@ -378,7 +378,7 @@ mod tests {
|
||||
fs::write(dir.path().join("tokenizer.json"), "{}").expect("write tokenizer");
|
||||
fs::write(
|
||||
dir.path().join("tokenizer_config.json"),
|
||||
r#"{"tokenizer_class":"PreTrainedTokenizerFast"}"#,
|
||||
r#"{"tokenizer_class":"TokenizersBackend"}"#,
|
||||
)
|
||||
.expect("write tokenizer config");
|
||||
fs::write(dir.path().join("config.json"), "{}").expect("write config");
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user