Compare commits

..
Author SHA1 Message Date
Bugen ZhaoandOpenAI Codex e1a763558c [CI] Discover Rust coverage artifacts from build metadata
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-23 06:44:16 +00:00
Bugen ZhaoandOpenAI Codex 84aeec9f22 [CI] Simplify Rust coverage reporting
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 08:18:14 +00:00
Bugen ZhaoandOpenAI Codex 82a770ddbd [CI] Simplify Rust coverage aggregation
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-22 02:58:45 +00:00
Bugen Zhao a09a9bace1 [CI] Disable redundant Codecov file fixes 2026-07-21 13:57:38 +00:00
Bugen Zhao 6c20d467a2 [CI] Run Codecov from repository root 2026-07-21 13:34:22 +00:00
Bugen ZhaoandOpenAI Codex cb59d0a351 [CI] Collect Rust coverage in Buildkite
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 13:01:56 +00:00
Bugen ZhaoandOpenAI Codex 0ab1bded36 [CI] Instrument Rust artifacts for coverage
Co-authored-by: OpenAI Codex <codex@openai.com>
2026-07-21 12:22:25 +00:00
213 changed files with 3121 additions and 5937 deletions
+28 -28
View File
@@ -17,7 +17,7 @@ DEFAULT_REPO_SLUG="vllm-project/vllm"
DEFAULT_CI_HCL_SOURCE="docker/ci-rocm.hcl"
DEFAULT_CI_BASE_CONTENT_FILES="requirements/common.txt requirements/rocm.txt requirements/test/rocm.txt docker/Dockerfile.rocm_base docker/ci-rocm.hcl docker/docker-bake-rocm.hcl tools/install_torchcodec_rocm.sh tools/install_protoc.sh rust-toolchain.toml tests/vllm_test_utils .buildkite/scripts/ci-bake-rocm.sh .buildkite/scripts/rocm/build-ci-base.sh"
DEFAULT_CI_BASE_DOCKERFILE="docker/Dockerfile.rocm"
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_nixl build_rocshmem build_deepep mori_base ci_base"
DEFAULT_CI_BASE_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust_toolchain_input_1 rust-toolchain-input rust-toolchain build_rixl build_rocshmem build_deepep mori_base ci_base"
DEFAULT_CI_BASE_METADATA_VERSION="1"
IMAGE_EXISTED_BEFORE_BUILD=0
@@ -1159,8 +1159,8 @@ ci_base_metadata_pairs() {
metadata_pair "vllm.rocm.nic_backend" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIC_BACKEND")"
metadata_pair "vllm.rocm.ainic_version" "$(resolve_dockerfile_arg_value "${dockerfile}" "AINIC_VERSION")"
metadata_pair "vllm.rocm.ubuntu_codename" "$(resolve_dockerfile_arg_value "${dockerfile}" "UBUNTU_CODENAME")"
metadata_pair "vllm.rocm.nixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_REPO")"
metadata_pair "vllm.rocm.nixl_commit" "${NIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "NIXL_BRANCH")}"
metadata_pair "vllm.rocm.rixl_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_REPO")"
metadata_pair "vllm.rocm.rixl_commit" "${RIXL_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "RIXL_BRANCH")}"
metadata_pair "vllm.rocm.ucx_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_REPO")"
metadata_pair "vllm.rocm.ucx_commit" "${UCX_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "UCX_BRANCH")}"
metadata_pair "vllm.rocm.rocshmem_repo" "$(resolve_dockerfile_arg_value "${dockerfile}" "ROCSHMEM_REPO")"
@@ -1169,7 +1169,7 @@ ci_base_metadata_pairs() {
metadata_pair "vllm.rocm.deepep_commit" "${DEEPEP_BRANCH:-$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_BRANCH")}"
metadata_pair "vllm.rocm.deepep_nic" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_NIC")"
metadata_pair "vllm.rocm.deepep_rocm_arch" "$(resolve_dockerfile_arg_value "${dockerfile}" "DEEPEP_ROCM_ARCH")"
metadata_pair "vllm.rocm.nixl_cache_key" "${NIXL_CACHE_KEY:-}"
metadata_pair "vllm.rocm.rixl_cache_key" "${RIXL_CACHE_KEY:-}"
metadata_pair "vllm.rocm.rocshmem_cache_key" "${ROCSHMEM_CACHE_KEY:-}"
metadata_pair "vllm.rocm.deepep_cache_key" "${DEEPEP_CACHE_KEY:-}"
@@ -1686,7 +1686,7 @@ extract_dependency_pins() {
return 0
fi
for var in NIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
for var in RIXL_BRANCH UCX_BRANCH ROCSHMEM_BRANCH DEEPEP_BRANCH; do
if [[ -n "${!var:-}" ]]; then
echo "Using provided ${var}: ${!var}"
continue
@@ -1706,30 +1706,30 @@ extract_dependency_pins() {
compute_dependency_cache_keys() {
local bake_dir=""
local dockerfile_rocm=""
local nixl_branch=""
local rixl_branch=""
local ucx_branch=""
local rocshmem_branch=""
local deepep_branch=""
local nixl_material=""
local rixl_material=""
local rocshmem_material=""
local deepep_material=""
bake_dir=$(dirname "${VLLM_BAKE_FILE}")
dockerfile_rocm="${bake_dir}/Dockerfile.rocm"
nixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "NIXL_BRANCH")
rixl_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "RIXL_BRANCH")
ucx_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "UCX_BRANCH")
rocshmem_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "ROCSHMEM_BRANCH")
deepep_branch=$(resolve_dockerfile_arg_value "${dockerfile_rocm}" "DEEPEP_BRANCH")
if [[ -n "${nixl_branch}" && -n "${ucx_branch}" ]]; then
nixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_nixl")
NIXL_CACHE_KEY=$(
if [[ -n "${rixl_branch}" && -n "${ucx_branch}" ]]; then
rixl_material=$(compose_stage_cache_material "${dockerfile_rocm}" "base build_rixl")
RIXL_CACHE_KEY=$(
compose_dependency_cache_key \
"${nixl_branch}-ucx-${ucx_branch}" \
"${nixl_material}"
"${rixl_branch}-ucx-${ucx_branch}" \
"${rixl_material}"
)
export NIXL_CACHE_KEY
echo "NIXL dependency cache key: ${NIXL_CACHE_KEY}"
export RIXL_CACHE_KEY
echo "RIXL dependency cache key: ${RIXL_CACHE_KEY}"
fi
if [[ -n "${rocshmem_branch}" ]]; then
@@ -1780,11 +1780,11 @@ dependency_cache_ref_for_target() {
local cache_repo="${DOCKERHUB_CACHE_REPO:-rocm/vllm-ci-cache}"
case "${target}" in
nixl-rocm-ci)
if [[ -n "${NIXL_CACHE_KEY:-}" ]]; then
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_CACHE_KEY}"
elif [[ -n "${NIXL_BRANCH:-}" ]]; then
printf '%s\n' "${cache_repo}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
rixl-rocm-ci)
if [[ -n "${RIXL_CACHE_KEY:-}" ]]; then
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_CACHE_KEY}"
elif [[ -n "${RIXL_BRANCH:-}" ]]; then
printf '%s\n' "${cache_repo}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH:-}"
fi
;;
rocshmem-rocm-ci)
@@ -1815,7 +1815,7 @@ add_dependency_cache_target() {
resolve_ci_base_dependency_targets() {
local mode="${ROCM_DEP_CACHE_EXPORT_MODE:-missing}"
local nixl_ref=""
local rixl_ref=""
local rocshmem_ref=""
local deepep_ref=""
@@ -1824,7 +1824,7 @@ resolve_ci_base_dependency_targets() {
case "${mode}" in
always)
echo "ROCM_DEP_CACHE_EXPORT_MODE=always; exporting all dependency caches serially"
for target in nixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
for target in rixl-rocm-ci rocshmem-rocm-ci deepep-rocm-ci; do
if [[ -n "$(dependency_cache_ref_for_target "${target}")" ]]; then
add_dependency_cache_target "${target}"
fi
@@ -1844,13 +1844,13 @@ resolve_ci_base_dependency_targets() {
;;
esac
if [[ "${mode}" != "always" && -n "${NIXL_CACHE_KEY:-}" ]]; then
nixl_ref=$(dependency_cache_ref_for_target "nixl-rocm-ci")
if dependency_cache_ref_exists "${nixl_ref}"; then
echo "NIXL dependency cache exists: ${nixl_ref}"
if [[ "${mode}" != "always" && -n "${RIXL_CACHE_KEY:-}" ]]; then
rixl_ref=$(dependency_cache_ref_for_target "rixl-rocm-ci")
if dependency_cache_ref_exists "${rixl_ref}"; then
echo "RIXL dependency cache exists: ${rixl_ref}"
else
echo "NIXL dependency cache missing; will seed: ${nixl_ref}"
add_dependency_cache_target "nixl-rocm-ci"
echo "RIXL dependency cache missing; will seed: ${rixl_ref}"
add_dependency_cache_target "rixl-rocm-ci"
fi
fi
@@ -29,6 +29,7 @@ PYO3_PYTHON_VERSION="${PYO3_PYTHON_VERSION:-3.12}"
CARGO_SORT_VERSION_REQ="${CARGO_SORT_VERSION_REQ:-2}"
CARGO_DENY_VERSION_REQ="${CARGO_DENY_VERSION_REQ:-0.20}"
CARGO_NEXTEST_VERSION_REQ="${CARGO_NEXTEST_VERSION_REQ:-0.9}"
CARGO_LLVM_COV_VERSION="${CARGO_LLVM_COV_VERSION:-0.8.7}"
log_section() {
echo "--- $*"
@@ -106,6 +107,18 @@ install_cargo_nextest() {
"cargo-nextest@${CARGO_NEXTEST_VERSION_REQ}"
}
install_cargo_llvm_cov() {
log_section "Installing cargo-llvm-cov ${CARGO_LLVM_COV_VERSION}"
local toolchain
toolchain="$(rust_toolchain)"
rustup component add --toolchain "$toolchain" llvm-tools-preview
cargo binstall \
--no-confirm \
--force \
--secure \
"cargo-llvm-cov@${CARGO_LLVM_COV_VERSION}"
}
install_uv() {
log_section "Installing uv ${UV_VERSION}"
curl -L --proto '=https' --tlsv1.2 -sSf \
@@ -176,14 +189,41 @@ run_tests() {
setup_pyo3_python
install_cargo_binstall
install_cargo_nextest
install_cargo_llvm_cov
log_section "Running cargo nextest"
cargo nextest run \
log_section "Running cargo nextest with Rust coverage"
mkdir -p artifacts
export LLVM_PROFILE_FILE_NAME="vllm-rust-unit-%4m.profraw"
cargo llvm-cov clean \
--manifest-path rust/Cargo.toml \
--profraw-only
set +e
cargo llvm-cov nextest \
--manifest-path rust/Cargo.toml \
--workspace \
--all-features \
--locked \
--no-fail-fast
--no-fail-fast \
--no-clean \
--lcov \
--output-path artifacts/rust-unit.lcov \
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/'
local coverage_rc=$?
local upload_rc=0
if [[ $coverage_rc -eq 0 ]]; then
# shellcheck source=.buildkite/scripts/rust-coverage.sh
source .buildkite/scripts/rust-coverage.sh
rust_coverage_upload artifacts/rust-unit.lcov rust-unit
upload_rc=$?
fi
set -e
if [[ $coverage_rc -ne 0 ]]; then
return "$coverage_rc"
fi
return "$upload_rc"
}
install_protoc
+182
View File
@@ -0,0 +1,182 @@
#!/bin/sh
RUST_CODECOV_VERSION="v11.3.1"
RUST_CODECOV_SHA256="ca1d64196d2d34771084afe76ea657d581bf628e31d993ff8e52ea09cc88a56d"
rust_coverage_repo_root() {
if [ -f /vllm-workspace/.buildkite/scripts/rust-coverage.sh ]; then
printf '%s\n' /vllm-workspace
elif [ -n "${BUILDKITE_BUILD_CHECKOUT_PATH:-}" ] \
&& [ -d "$BUILDKITE_BUILD_CHECKOUT_PATH" ]; then
printf '%s\n' "$BUILDKITE_BUILD_CHECKOUT_PATH"
else
git rev-parse --show-toplevel
fi
}
rust_coverage_start() {
RUST_COVERAGE_FLAG=${1:?coverage flag is required}
RUST_COVERAGE_DIR="/tmp/vllm-rust-coverage/${BUILDKITE_JOB_ID:-local}"
export RUST_COVERAGE_FLAG RUST_COVERAGE_DIR
mkdir -p "$RUST_COVERAGE_DIR"
LLVM_PROFILE_FILE="$RUST_COVERAGE_DIR/rust-%4m.profraw"
export LLVM_PROFILE_FILE
trap rust_coverage_finalize 0
}
rust_coverage_objects() {
rust_cov_objects_manifest="$(dirname "$(command -v llvm-cov)")/../objects"
python3 - "$rust_cov_objects_manifest" <<'PY'
from pathlib import Path
import sys
for relative in Path(sys.argv[1]).read_text().splitlines():
for entry in sys.path:
path = Path(entry or ".").resolve() / relative
if path.is_file():
print(path)
break
else:
raise RuntimeError(f"installed Rust coverage object was not found: {relative}")
PY
}
rust_coverage_collect() {
rust_cov_collect_flag=${1:?coverage flag is required}
rust_cov_collect_lcov="$RUST_COVERAGE_DIR/$rust_cov_collect_flag.lcov"
rust_cov_collect_objects=$(rust_coverage_objects) || return 1
rust_cov_collect_primary=
set --
while IFS= read -r rust_cov_collect_object; do
if [ -z "$rust_cov_collect_primary" ]; then
rust_cov_collect_primary=$rust_cov_collect_object
else
set -- "$@" "--object=$rust_cov_collect_object"
fi
done <<EOF
$rust_cov_collect_objects
EOF
llvm-profdata merge \
-sparse \
"$RUST_COVERAGE_DIR"/*.profraw \
-o "$RUST_COVERAGE_DIR/merged.profdata" || return 1
llvm-cov export \
"$rust_cov_collect_primary" \
"$@" \
--format=lcov \
--instr-profile="$RUST_COVERAGE_DIR/merged.profdata" \
--ignore-filename-regex='/\.cargo/(registry|git)/|/rustc/|/target/' \
> "$rust_cov_collect_lcov" || return 1
RUST_COVERAGE_LCOV=$rust_cov_collect_lcov
export RUST_COVERAGE_LCOV
}
rust_coverage_upload() {
rust_cov_upload_lcov=${1:?LCOV path is required}
rust_cov_upload_flag=${2:?coverage flag is required}
rust_cov_upload_repo_root=$(rust_coverage_repo_root) || return 1
if [ "$(uname -m)" != "x86_64" ]; then
echo "Rust coverage upload currently supports x86_64 CI agents" >&2
return 1
fi
rust_cov_upload_codecov_dir=$(mktemp -d /tmp/codecov-bin.XXXXXX) \
|| return 1
curl -fsSL \
"https://github.com/codecov/codecov-cli/releases/download/${RUST_CODECOV_VERSION}/codecovcli_linux" \
-o "$rust_cov_upload_codecov_dir/codecov" || return 1
echo "$RUST_CODECOV_SHA256 $rust_cov_upload_codecov_dir/codecov" \
| sha256sum -c - || return 1
chmod +x "$rust_cov_upload_codecov_dir/codecov" || return 1
rust_cov_upload_slug="vllm-project/vllm"
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ] \
&& [ -n "${BUILDKITE_PULL_REQUEST_REPO:-}" ]; then
rust_cov_upload_slug=$(echo "$BUILDKITE_PULL_REQUEST_REPO" \
| sed -E 's#(git@|https?://)([^/:]+)[:/]([^/]+/[^/.]+)(\.git)?$#\3#')
case "$rust_cov_upload_slug" in
*/*) ;;
*) rust_cov_upload_slug="vllm-project/vllm" ;;
esac
fi
rust_cov_upload_branch=${BUILDKITE_BRANCH:?BUILDKITE_BRANCH is required}
if [ -z "${CODECOV_TOKEN:-}" ]; then
# Codecov accepts tokenless public uploads on unprotected branch names.
# A colon-separated prefix keeps feature-branch and fork uploads from
# requiring a repository secret.
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
rust_cov_upload_branch="pr${BUILDKITE_PULL_REQUEST}:$rust_cov_upload_branch"
else
rust_cov_upload_branch="buildkite:$rust_cov_upload_branch"
fi
fi
set --
set -- "$@" upload-process
set -- "$@" --file "$rust_cov_upload_lcov"
# LCOV paths are mapped server-side by codecov.yml. Skip the CLI's local
# source-line fix scanning, which is unrelated to path mapping.
set -- "$@" --disable-search --disable-file-fixes
set -- "$@" --fail-on-error --git-service github
set -- "$@" --build "${BUILDKITE_BUILD_NUMBER:?BUILDKITE_BUILD_NUMBER is required}"
set -- "$@" --branch "$rust_cov_upload_branch"
set -- "$@" --sha "${BUILDKITE_COMMIT:?BUILDKITE_COMMIT is required}"
set -- "$@" --slug "$rust_cov_upload_slug"
set -- "$@" --flag "$rust_cov_upload_flag"
set -- "$@" --name "${rust_cov_upload_flag}-${BUILDKITE_JOB_ID:?BUILDKITE_JOB_ID is required}"
set -- "$@" --dir "$rust_cov_upload_repo_root"
set -- "$@" --network-root-folder "$rust_cov_upload_repo_root"
if [ -n "${BUILDKITE_PULL_REQUEST:-}" ] \
&& [ "${BUILDKITE_PULL_REQUEST}" != "false" ]; then
set -- "$@" --pr "$BUILDKITE_PULL_REQUEST"
fi
rust_cov_upload_log="$rust_cov_upload_codecov_dir/codecov.log"
# E2E steps run from tests/, so execute from the repository root to resolve
# codecov.yml and repository paths consistently.
(
cd "$rust_cov_upload_repo_root" || exit 1
"$rust_cov_upload_codecov_dir/codecov" "$@"
) >"$rust_cov_upload_log" 2>&1
rust_cov_upload_rc=$?
cat "$rust_cov_upload_log"
# v11.3.1 can log API failures while returning zero even with
# --fail-on-error. Preserve the strict CI contract explicitly.
if grep -aEq 'error.* -- ' "$rust_cov_upload_log"; then
echo "Codecov CLI reported an upload error" >&2
rust_cov_upload_rc=1
fi
rm -rf "$rust_cov_upload_codecov_dir"
return "$rust_cov_upload_rc"
}
rust_coverage_finalize() {
rust_cov_finalize_test_rc=$?
trap - 0
set +e
rust_coverage_collect "$RUST_COVERAGE_FLAG"
rust_cov_finalize_collect_rc=$?
rust_cov_finalize_upload_rc=0
if [ "$rust_cov_finalize_collect_rc" -eq 0 ]; then
rust_coverage_upload "$RUST_COVERAGE_LCOV" "$RUST_COVERAGE_FLAG"
rust_cov_finalize_upload_rc=$?
fi
find "$RUST_COVERAGE_DIR" -type f -name '*.profraw' -delete
if [ "$rust_cov_finalize_test_rc" -ne 0 ]; then
exit "$rust_cov_finalize_test_rc"
fi
if [ "$rust_cov_finalize_collect_rc" -ne 0 ]; then
exit "$rust_cov_finalize_collect_rc"
fi
exit "$rust_cov_finalize_upload_rc"
}
+1 -1
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: Basic Correctness
key: basic-correctness
timeout_in_minutes: 68
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
+1 -1
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: Benchmarks CLI Test
key: benchmarks-cli-test
timeout_in_minutes: 45
timeout_in_minutes: 30
device: h200_18gb
source_file_dependencies:
- vllm/
-3
View File
@@ -26,10 +26,7 @@ steps:
- vllm/v1/cudagraph_dispatcher.py
- vllm/config/compilation.py
- vllm/compilation
- vllm/v1/worker/encoder_cudagraph.py
- vllm/v1/worker/encoder_cudagraph_defs.py
commands:
- pytest -v -s v1/cudagraph/test_cudagraph_dispatch.py
- pytest -v -s v1/cudagraph/test_cudagraph_mode.py
- pytest -v -s v1/cudagraph/test_breakable_cudagraph.py
- pytest -v -s v1/cudagraph/test_encoder_cudagraph.py
+1 -1
View File
@@ -51,7 +51,7 @@ steps:
- label: e2e Scheduling (1 GPU)
key: e2e-scheduling-1-gpu
timeout_in_minutes: 53
timeout_in_minutes: 35
device: h200_18gb
source_file_dependencies:
- vllm/v1/
+4 -4
View File
@@ -39,7 +39,7 @@ steps:
- label: Entrypoints Integration (API Server)
key: entrypoints-integration-api-server
device: h200_35gb
timeout_in_minutes: 75
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -59,7 +59,7 @@ steps:
- label: Entrypoints Integration (API Server OpenAI - Part 1)
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-1
timeout_in_minutes: 68
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -78,7 +78,7 @@ steps:
- label: Entrypoints Integration (API Server OpenAI - Part 2)
device: h200_35gb
key: entrypoints-integration-api-server-openai-part-2
timeout_in_minutes: 83
timeout_in_minutes: 45
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
@@ -156,7 +156,7 @@ steps:
- label: Entrypoints Integration (Pooling)
device: h200_35gb
key: entrypoints-integration-pooling
timeout_in_minutes: 75
timeout_in_minutes: 50
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/
+1 -4
View File
@@ -31,7 +31,7 @@ steps:
- label: V1 Sample + Logits
key: v1-sample-logits
timeout_in_minutes: 83
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/config/
@@ -90,7 +90,6 @@ steps:
- tests/v1/kv_offload
- tests/v1/simple_kv_offload
- tests/v1/worker
- tests/v1/streaming_input
- tests/v1/kv_connector/unit
- tests/v1/ec_connector/unit
- tests/v1/metrics
@@ -104,7 +103,6 @@ steps:
- pytest -v -s v1/kv_offload
- pytest -v -s v1/simple_kv_offload
- pytest -v -s v1/worker
- pytest -v -s v1/streaming_input
- pytest -v -s -m 'not cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'not cpu_test' v1/ec_connector/unit
- pytest -v -s -m 'not cpu_test' v1/metrics
@@ -145,7 +143,6 @@ steps:
- pytest -v -s -m 'cpu_test' v1/core
- pytest -v -s v1/structured_output
- pytest -v -s v1/test_serial_utils.py
- pytest -v -s v1/cudagraph/test_cudagraph_manager.py
- pytest -v -s -m 'cpu_test' v1/kv_connector/unit
- pytest -v -s -m 'cpu_test' v1/metrics
+1 -1
View File
@@ -5,7 +5,7 @@ steps:
- label: Model Executor
device: h200_35gb
key: model-executor
timeout_in_minutes: 60
timeout_in_minutes: 45
source_file_dependencies:
- vllm/engine/arg_utils.py
- vllm/config/model.py
-13
View File
@@ -46,19 +46,6 @@ steps:
depends_on:
- image-build-amd
- label: Inkling Unit Tests (B200)
key: inkling-unit-tests-b200
timeout_in_minutes: 40
device: b200-k8s
source_file_dependencies:
- vllm/models/inkling/
- vllm/cute_utils/
- cmake/external_projects/tml_fa4.cmake
- tests/models/inkling/
commands:
# FA4 kernel tests require SM100; the suite skips them elsewhere.
- pytest -v -s models/inkling
- label: Basic Models Test (Other CPU) # 5min
key: basic-models-test-other-cpu
depends_on:
+1 -1
View File
@@ -137,7 +137,7 @@ steps:
- label: Language Models Test (MTEB)
key: language-models-test-mteb
timeout_in_minutes: 68
timeout_in_minutes: 45
device: h200_18gb
optional: true
source_file_dependencies:
+4 -4
View File
@@ -4,7 +4,7 @@ depends_on:
steps:
- label: "Multi-Modal Models (Standard) 1: qwen2"
key: multi-modal-models-standard-1-qwen2
timeout_in_minutes: 68
timeout_in_minutes: 45
device: h200_18gb
source_file_dependencies:
- vllm/
@@ -20,7 +20,7 @@ steps:
- label: "Multi-Modal Models (Standard) 2: qwen3 + gemma"
key: multi-modal-models-standard-2-qwen3-gemma
timeout_in_minutes: 75
timeout_in_minutes: 50
device: h200_18gb
source_file_dependencies:
- vllm/
@@ -54,7 +54,7 @@ steps:
- label: "Multi-Modal Models (Standard) 4: other + whisper"
device: h200_35gb
key: multi-modal-models-standard-4-other-whisper
timeout_in_minutes: 75
timeout_in_minutes: 50
source_file_dependencies:
- vllm/
- tests/models/multimodal
@@ -85,7 +85,7 @@ steps:
- label: Multi-Modal Processor # 44min
key: multi-modal-processor
timeout_in_minutes: 98
timeout_in_minutes: 65
device: h200_18gb
source_file_dependencies:
- vllm/
+1 -1
View File
@@ -5,7 +5,7 @@ steps:
- label: PyTorch Compilation Unit Tests
device: h200_35gb
key: pytorch-compilation-unit-tests
timeout_in_minutes: 110
timeout_in_minutes: 90
source_file_dependencies:
- vllm/__init__.py
- vllm/_aiter_ops.py
+30
View File
@@ -8,6 +8,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/benchmarks/
- vllm/entrypoints/openai/
- vllm/entrypoints/serve/
@@ -23,6 +28,7 @@ steps:
- tests/entrypoints/openai/test_uds.py
- tests/v1/sample/test_logprobs_e2e.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- 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)"
@@ -43,6 +49,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- vllm/entrypoints/serve/
- vllm/v1/engine/
@@ -54,6 +65,7 @@ steps:
# - tests/entrypoints/serve/dev/test_sleep.py
- tests/entrypoints/serve/tokenize/test_tokenization.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- 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
@@ -72,10 +84,16 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- tests/utils.py
- tests/entrypoints/openai/correctness/test_lmeval.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- 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
@@ -86,11 +104,17 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/entrypoints/openai/
- vllm/tool_parsers/
- tests/utils.py
- tests/tool_use/
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- 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"
@@ -101,6 +125,11 @@ steps:
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- rust/
- build_rust.sh
- tools/build_rust.py
- rust-toolchain.toml
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
- vllm/distributed/
- vllm/engine/
- vllm/executor/
@@ -111,6 +140,7 @@ steps:
- tests/v1/distributed/test_hybrid_lb_dp.py
- tests/v1/distributed/test_internal_lb_dp.py
commands:
- . /vllm-workspace/.buildkite/scripts/rust-coverage.sh && rust_coverage_start rust-e2e
- export VLLM_USE_RUST_FRONTEND=1
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export NCCL_CUMEM_HOST_ENABLE=0
@@ -26,5 +26,7 @@ steps:
- rust-toolchain.toml
- .buildkite/test_areas/rust_frontend_cargo.yaml
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh
- .buildkite/scripts/rust-coverage.sh
- codecov.yml
commands:
- .buildkite/scripts/run-rust-frontend-cargo-ci.sh test
-16
View File
@@ -170,19 +170,3 @@ steps:
- tests/v1/e2e/spec_decode/
commands:
- pytest -v -s v1/e2e/spec_decode -k "qwen3_5-hybrid"
- label: Spec Decode DeepSeek MTP Parallel Load (B200)
key: spec-decode-deepseek-mtp-parallel-load-b200
timeout_in_minutes: 30
device: b200-k8s
optional: true
num_devices: 2
source_file_dependencies:
- vllm/v1/spec_decode/llm_base_proposer.py
- vllm/v1/spec_decode/eagle.py
- vllm/v1/worker/gpu/spec_decode/eagle/
- vllm/model_executor/models/deepseek_mtp.py
- vllm/model_executor/models/deepseek_v2.py
- tests/v1/e2e/spec_decode/test_mtp_parallel_load.py
commands:
- pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py
-12
View File
@@ -181,18 +181,6 @@ pull_request_rules:
add:
- performance
- name: label-quantization
description: Automatically apply quantization label
conditions:
- label != stale
- or:
- files~=^vllm/model_executor/layers/quantization/
- title~=(?i)quant
actions:
label:
add:
- quantization
- name: label-qwen
description: Automatically apply qwen label
conditions:
+1 -42
View File
@@ -130,47 +130,6 @@ jobs:
},
],
},
quantization: {
keywords: [
{
term: "quantization",
searchIn: "both"
},
{
term: "quantized",
searchIn: "both"
},
],
},
"intel-gpu": {
// Keyword search - matches whole words only (with word boundaries)
keywords: [
{
term: "B50",
searchIn: "both"
},
{
term: "B60",
searchIn: "both"
},
{
term: "B70",
searchIn: "both"
},
{
term: "intel gpu",
searchIn: "both"
},
{
term: "Arc GPU",
searchIn: "both"
},
{
term: "BMG",
searchIn: "both"
},
],
},
// Add more label configurations here as needed
// example: {
// keywords: [...],
@@ -532,4 +491,4 @@ jobs:
issue_number: context.issue.number,
body: message,
});
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
core.notice(`Requested missing ROCm info from @${author}: ${missing.map(m => m.name).join(', ')}`);
+1
View File
@@ -257,3 +257,4 @@ vllm/grpc/vllm_engine_pb2.pyi
# Ignore generated cpu headers
csrc/cpu/cpu_attn_dispatch_generated.h
rust-coverage-tools/
+37
View File
@@ -8,6 +8,8 @@
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "$0")" && pwd)"
CARGO_LLVM_COV_VERSION="0.8.7"
COVERAGE_TOOLS_DIR="$REPO_ROOT/rust-coverage-tools"
# Read the required toolchain from rust-toolchain.toml.
TOOLCHAIN=$(grep '^channel' "$REPO_ROOT/rust-toolchain.toml" | sed 's/.*= *"\(.*\)"/\1/')
@@ -30,4 +32,39 @@ else
PROFILE_ARG="--release"
fi
rm -rf "$COVERAGE_TOOLS_DIR"
mkdir -p "$COVERAGE_TOOLS_DIR/bin" "$COVERAGE_TOOLS_DIR/lib"
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
# rustc wrapper flags are invisible to Cargo's normal fingerprinting.
# Keep instrumented intermediates isolated when local builds switch modes.
export CARGO_TARGET_DIR="$REPO_ROOT/rust/target/coverage"
rustup component add --toolchain "$TOOLCHAIN" llvm-tools-preview
cargo +"$TOOLCHAIN" install \
--locked \
--version "$CARGO_LLVM_COV_VERSION" \
cargo-llvm-cov
eval "$(
cargo +"$TOOLCHAIN" llvm-cov show-env \
--manifest-path "$REPO_ROOT/rust/Cargo.toml" \
--sh
)"
# Build scripts and proc macros can run during compilation. Their profiles
# are unrelated to runtime coverage and would otherwise pollute the tree.
export LLVM_PROFILE_FILE=/dev/null
export VLLM_RUST_COVERAGE_OBJECTS="$COVERAGE_TOOLS_DIR/objects"
fi
python3 "$REPO_ROOT/tools/build_rust.py" "$PROFILE_ARG"
if [[ "${VLLM_RUST_COVERAGE:-0}" == "1" ]]; then
LLVM_BIN_DIR="$(dirname "$(rustup run "$TOOLCHAIN" rustc \
--print target-libdir)")/bin"
cp "$LLVM_BIN_DIR"/{llvm-cov,llvm-profdata} "$COVERAGE_TOOLS_DIR/bin/"
chmod 0755 "$COVERAGE_TOOLS_DIR/bin/"*
cp -L "$LLVM_BIN_DIR"/../lib/libLLVM.so* "$COVERAGE_TOOLS_DIR/lib/"
chmod 0644 "$COVERAGE_TOOLS_DIR/lib/"*
fi
+1 -1
View File
@@ -17,7 +17,7 @@ else()
FetchContent_Declare(
fmha_sm100
GIT_REPOSITORY https://github.com/vllm-project/MSA.git
GIT_TAG 890aaa1a37a598ad17ccff0827fea21540d381fa
GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73
GIT_PROGRESS TRUE
CONFIGURE_COMMAND ""
BUILD_COMMAND ""
@@ -39,7 +39,7 @@ else()
FetchContent_Declare(
vllm-flash-attn
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
GIT_TAG ed4b7342bc8f0489dd9b649d5288867e35fc6a32
GIT_TAG 168920233059c48de6199e2cda74003b2ce3d199
GIT_PROGRESS TRUE
# Don't share the vllm-flash-attn build between build types
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
+13
View File
@@ -10,3 +10,16 @@ fixes:
- "/usr/local/lib/python3.*/site-packages/vllm/::vllm/"
- "/usr/lib/python3.*/dist-packages/vllm/::vllm/"
- "/usr/lib/python3.*/site-packages/vllm/::vllm/"
# Map Rust sources built in the E2E image and on Buildkite agents.
- "/workspace/rust/::rust/"
- "/var/lib/buildkite-agent/.*/rust/::rust/"
flags:
rust-unit:
paths:
- rust/
carryforward: false
rust-e2e:
paths:
- rust/
carryforward: false
-3
View File
@@ -1025,9 +1025,6 @@ __global__ void gather_and_maybe_dequant_cache(
batch_offset += offset;
int32_t block_table_id = batch_offset / block_size;
int32_t slot_id = batch_offset % block_size;
// seq_starts may push the block index past the end of the batch's block
// table row.
if (block_table_id >= block_table_stride) continue;
int32_t block_table_offset = batch_id * block_table_stride + block_table_id;
int32_t block_id = block_table[block_table_offset];
int64_t cache_offset =
+3 -6
View File
@@ -9,16 +9,14 @@ void topk_softmax(torch::stable::Tensor& topk_weights,
torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& gating_output, bool renormalize,
std::optional<torch::stable::Tensor> bias,
std::optional<torch::stable::Tensor> is_padding);
std::optional<torch::stable::Tensor> bias);
void topk_sigmoid(torch::stable::Tensor& topk_weights,
torch::stable::Tensor& topk_indices,
torch::stable::Tensor& token_expert_indices,
torch::stable::Tensor& gating_output, bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor,
std::optional<torch::stable::Tensor> is_padding);
double routed_scaling_factor);
void topk_softplus_sqrt(
torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices,
@@ -27,8 +25,7 @@ void topk_softplus_sqrt(
double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& is_padding);
const std::optional<torch::stable::Tensor>& tid2eid);
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
std::optional<torch::stable::Tensor> topk_ids,
@@ -174,8 +174,7 @@ __launch_bounds__(TPB) __global__ void moeTopK(
const int end_expert,
const bool renormalize,
const float* bias,
const double routed_scaling_factor,
const bool* is_padding)
const double routed_scaling_factor)
{
using cub_kvp = cub::KeyValuePair<int, float>;
@@ -229,14 +228,12 @@ __launch_bounds__(TPB) __global__ void moeTopK(
const int expert = result_kvp.key;
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const bool is_pad_row = is_padding != nullptr && is_padding[block_row];
const int idx = k * block_row + k_idx;
// Return the unbiased scores for output weights
output[idx] = inputs_after_softmax[thread_read_offset + expert];
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
: (should_process_row ? (expert - start_expert) : num_experts);
assert(is_pad_row || indices[idx] >= 0);
indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
assert(indices[idx] >= 0);
source_rows[idx] = k_idx * num_rows + block_row;
if (renormalize) {
selected_sum += inputs_after_softmax[thread_read_offset + expert];
@@ -280,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 double routed_scaling_factor, const bool* is_padding)
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>,
@@ -548,14 +545,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
// Add a guard to ignore experts not included by this node
const bool node_uses_expert = expert >= start_expert && expert < end_expert;
const bool should_process_row = row_is_active && node_uses_expert;
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
// The lead thread from each sub-group will write out the final results to global memory. (This will be a
// single) thread per row of the input/output matrices.
const int idx = k * thread_row + k_idx;
output[idx] = max_val;
indices[idx] = is_pad_row ? static_cast<IndType>(-1)
: (should_process_row ? (expert - start_expert) : NUM_EXPERTS);
indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
source_rows[idx] = k_idx * num_rows + thread_row;
if (renormalize) {
selected_sum += max_val;
@@ -610,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, const double routed_scaling_factor, cudaStream_t stream, const bool* is_padding)
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>;
@@ -621,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, routed_scaling_factor, is_padding);
input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor);
}
#ifndef USE_ROCM
@@ -632,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, routed_scaling_factor, stream, is_padding);
bias, routed_scaling_factor, stream);
#else
#define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
@@ -640,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, routed_scaling_factor, stream, is_padding); \
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, routed_scaling_factor, stream, is_padding); \
bias, routed_scaling_factor, stream); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
@@ -666,8 +661,7 @@ void topkGatingKernelLauncher(
const bool renormalize,
const float* bias,
const double routed_scaling_factor,
cudaStream_t stream,
const bool* is_padding) {
cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4;
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
#ifndef USE_ROCM
@@ -742,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, routed_scaling_factor, is_padding);
num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor);
}
}
}
@@ -761,8 +755,7 @@ void dispatch_topk_launch(
int num_tokens, int num_experts, int topk, bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor,
cudaStream_t stream,
std::optional<torch::stable::Tensor> is_padding)
cudaStream_t stream)
{
const float* bias_ptr = nullptr;
if (bias.has_value()) {
@@ -776,18 +769,6 @@ void dispatch_topk_launch(
bias_ptr = bias_tensor.const_data_ptr<float>();
}
const bool* is_padding_ptr = nullptr;
if (is_padding.has_value()) {
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
STD_TORCH_CHECK(is_padding_tensor.scalar_type() == torch::headeronly::ScalarType::Bool,
"is_padding tensor must be bool");
STD_TORCH_CHECK(is_padding_tensor.dim() == 1, "is_padding tensor must be 1D");
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
"is_padding size mismatch, expected: ", num_tokens);
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), "is_padding tensor must be contiguous");
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
}
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -796,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, routed_scaling_factor, stream, is_padding_ptr);
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()),
@@ -805,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, routed_scaling_factor, stream, is_padding_ptr);
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>(
@@ -815,7 +796,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
num_tokens, num_experts, topk, renormalize,
bias_ptr, routed_scaling_factor, stream, is_padding_ptr);
bias_ptr, routed_scaling_factor, stream);
}
}
@@ -825,8 +806,7 @@ void topk_softmax(
torch::stable::Tensor& token_expert_indices, // [num_tokens, topk]
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize,
std::optional<torch::stable::Tensor> bias,
std::optional<torch::stable::Tensor> is_padding)
std::optional<torch::stable::Tensor> bias)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
@@ -845,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, 1.0, stream, is_padding);
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, 1.0, stream, is_padding);
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, 1.0, stream, is_padding);
bias, 1.0, stream);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
@@ -866,8 +846,7 @@ void topk_sigmoid(
torch::stable::Tensor& gating_output, // [num_tokens, num_experts]
bool renormalize,
std::optional<torch::stable::Tensor> bias,
double routed_scaling_factor,
std::optional<torch::stable::Tensor> is_padding)
double routed_scaling_factor)
{
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
@@ -886,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, routed_scaling_factor, stream, is_padding);
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, routed_scaling_factor, stream, is_padding);
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, routed_scaling_factor, stream, is_padding);
bias, routed_scaling_factor, stream);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type());
}
@@ -80,27 +80,22 @@ __launch_bounds__(128) __global__
OutIndType* indices, int num_rows,
int num_experts, float routed_scaling_factor,
const HashIndType* input_ids,
const HashIndType* tid2eid,
const bool* is_padding) {
const HashIndType* tid2eid) {
const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int lane = threadIdx.x % 32;
if (warp >= num_rows) return;
const int64_t token_id = load_index_as_int64(input_ids, warp);
const bool is_pad_row = is_padding != nullptr && is_padding[warp];
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
cudaGridDependencySynchronize();
#endif
int expert = 0;
float weight = 0.f;
if (lane < 6 && !is_pad_row) {
if (lane < 6) {
// only load and calculate for 6 experts
expert = static_cast<int>(tid2eid[token_id * 6 + lane]);
const float x = input[warp * num_experts + expert];
weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x))));
if (isnan(weight)) {
weight = 0.f;
}
}
float weight_sum = weight;
#pragma unroll
@@ -116,8 +111,7 @@ __launch_bounds__(128) __global__
const int offset = warp * 6 + lane;
output[offset] =
weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f);
indices[offset] = !is_pad_row ? static_cast<OutIndType>(expert)
: static_cast<OutIndType>(-1);
indices[offset] = static_cast<OutIndType>(expert);
}
}
@@ -126,8 +120,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
int num_rows, int num_experts,
double routed_scaling_factor,
const HashIndType* input_ids,
const HashIndType* tid2eid, cudaStream_t stream,
const bool* is_padding) {
const HashIndType* tid2eid, cudaStream_t stream) {
if (num_rows == 0) return;
auto* kernel = &dsv4HashTopkSoftplusSqrt<OutIndType, HashIndType>;
cudaLaunchConfig_t config = {};
@@ -141,7 +134,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
config.numAttrs = 1;
const float scale = static_cast<float>(routed_scaling_factor);
cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows,
num_experts, scale, input_ids, tid2eid, is_padding);
num_experts, scale, input_ids, tid2eid);
}
#endif
@@ -173,8 +166,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
const int num_rows, IndType* indices, int* source_rows, const int k,
const int start_expert, const int end_expert, const bool renormalize,
double routed_scaling_factor, const float* correction_bias,
const HashIndType* input_ids, const HashIndType* tid2eid,
const bool* is_padding) {
const HashIndType* input_ids, const HashIndType* tid2eid) {
static_assert(std::is_same_v<InputType, float> ||
std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -239,7 +231,6 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
return;
}
const bool row_is_active = finished ? !finished[thread_row] : true;
const bool is_pad_row = is_padding != nullptr && is_padding[thread_row];
// We finally start setting up the read pointers for each thread. First, each
// thread jumps to the start of the row it will read.
@@ -258,12 +249,9 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
cudaGridDependencySynchronize();
#endif
if (is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = 0.f;
}
} else if constexpr (std::is_same_v<InputType, float>) {
// NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert
// to float
if constexpr (std::is_same_v<InputType, float>) {
using VecType = AlignedArray<float, ELTS_PER_LDG>;
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&row_chunk);
const VecType* vec_thread_read_ptr =
@@ -327,22 +315,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
if constexpr (USE_HASH) {
const int64_t token_id = load_index_as_int64(input_ids, thread_row);
const int64_t token_expert_offset = token_id * static_cast<int64_t>(k);
if (!is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
// Dummy/padding tokens can result in NaN values, so
// clamp them to 0.0. Note: this clamp could likely be removed if
// 'is_padding' is made mandatory
if (isnan(val)) {
val = 0.f;
}
row_chunk[ii] = val;
}
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
row_chunk[ii] = sqrtf(val);
}
float selected_sum = 0.f;
#pragma unroll
@@ -357,8 +335,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
if (expert == expert_idx) {
indices[idx] = !is_pad_row ? static_cast<IndType>(expert)
: static_cast<IndType>(-1);
indices[idx] = static_cast<IndType>(expert);
selected_sum += row_chunk[ii];
break;
}
@@ -402,31 +379,23 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
#endif
return;
} else {
if (!is_pad_row) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
// Compute softplus: log(1 + exp(val)) with numerical stability
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
// Dummy/padding tokens can result in NaN values, so
// clamp them to 0.0. Note: this clamp could likely be removed if
// 'is_padding' is made mandatory
if (isnan(val)) {
val = 0.f;
}
if (correction_bias) {
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
val = val + correction_bias[expert_idx];
}
row_chunk[ii] = val;
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
float val_b = val * beta;
// Compute softplus: log(1 + exp(val)) with numerical stability
// When val > threshold, softplus(x) ≈ x to avoid exp overflow
val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta;
val = sqrtf(val);
if (correction_bias) {
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread +
group_id * THREADS_PER_ROW * ELTS_PER_LDG +
local_id;
val = val + correction_bias[expert_idx];
}
row_chunk[ii] = val;
}
// Original TopK path: find top-k experts by score
@@ -481,19 +450,18 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
// Add a guard to ignore experts not included by this node
const bool node_uses_expert =
expert >= start_expert && expert < end_expert;
const bool should_process_row =
row_is_active && node_uses_expert && !is_pad_row;
const bool should_process_row = row_is_active && node_uses_expert;
// The lead thread from each sub-group will write out the final results
// to global memory. (This will be a single) thread per row of the
// input/output matrices.
const int idx = k * thread_row + k_idx;
if (correction_bias != nullptr && should_process_row) {
if (correction_bias != nullptr) {
max_val -= correction_bias[expert];
}
output[idx] = max_val;
indices[idx] =
!is_pad_row ? expert - start_expert : static_cast<IndType>(-1);
should_process_row ? (expert - start_expert) : NUM_EXPERTS;
source_rows[idx] = k_idx * num_rows + thread_row;
if (renormalize) {
selected_sum += max_val;
@@ -576,7 +544,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
const int start_expert, const int end_expert, const bool renormalize,
double routed_scaling_factor, const float* correction_bias,
const bool use_hash, const HashIndType* input_ids,
const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) {
const HashIndType* tid2eid, cudaStream_t stream) {
static constexpr int BYTES_PER_LDG =
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
using Constants =
@@ -605,12 +573,12 @@ void topkGatingSoftplusSqrtLauncherHelper(
cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows,
indices, source_row, k, start_expert, end_expert,
renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, is_padding);
input_ids, tid2eid);
#else
kernel<<<num_blocks, block_dim, 0, stream>>>(
input, finished, output, num_rows, indices, source_row, k, start_expert,
end_expert, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, is_padding);
input_ids, tid2eid);
#endif
})
}
@@ -624,7 +592,7 @@ void topkGatingSoftplusSqrtLauncherHelper(
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \
stream, is_padding);
stream);
#else
#define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \
if (WARP_SIZE == 64) { \
@@ -633,14 +601,14 @@ void topkGatingSoftplusSqrtLauncherHelper(
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, \
tid2eid, stream, is_padding); \
tid2eid, stream); \
} else if (WARP_SIZE == 32) { \
topkGatingSoftplusSqrtLauncherHelper<NUM_EXPERTS, WARPS_PER_TB, 32, \
MAX_BYTES>( \
gating_output, nullptr, topk_weights, topk_indices, \
token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \
routed_scaling_factor, correction_bias, use_hash, input_ids, \
tid2eid, stream, is_padding); \
tid2eid, stream); \
} else { \
assert(false && \
"Unsupported warp size. Only 32 and 64 are supported for ROCm"); \
@@ -654,14 +622,14 @@ void topkGatingSoftplusSqrtKernelLauncher(
const int topk, const bool renormalize, double routed_scaling_factor,
const float* correction_bias, const bool use_hash,
const HashIndType* input_ids, const HashIndType* tid2eid,
cudaStream_t stream, const bool* is_padding) {
cudaStream_t stream) {
#ifndef USE_ROCM
if constexpr (std::is_same_v<InputType, float>) {
if (use_hash && topk == 6 && renormalize &&
(num_experts == 256 || num_experts == 384)) {
launchDsv4HashTopk<IndType, HashIndType>(
gating_output, topk_weights, topk_indices, num_tokens, num_experts,
routed_scaling_factor, input_ids, tid2eid, stream, is_padding);
routed_scaling_factor, input_ids, tid2eid, stream);
return;
}
}
@@ -760,8 +728,7 @@ void dispatch_topk_softplus_sqrt_launch(
int num_experts, int topk, bool renormalize, double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,
const std::optional<torch::stable::Tensor>& is_padding) {
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
bias_ptr = correction_bias.value().const_data_ptr<float>();
@@ -770,22 +737,6 @@ void dispatch_topk_softplus_sqrt_launch(
auto launch = [&](auto* topk_indices_ptr) {
using OutIndType =
typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
const bool* is_padding_ptr = nullptr;
if (is_padding.has_value()) {
const torch::stable::Tensor& is_padding_tensor = is_padding.value();
STD_TORCH_CHECK(is_padding_tensor.scalar_type() ==
torch::headeronly::ScalarType::Bool,
"is_padding tensor must be bool");
STD_TORCH_CHECK(is_padding_tensor.dim() == 1,
"is_padding tensor must be 1D");
STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens,
"is_padding size mismatch, expected: ", num_tokens);
STD_TORCH_CHECK(is_padding_tensor.is_contiguous(),
"is_padding tensor must be contiguous");
is_padding_ptr = is_padding_tensor.const_data_ptr<bool>();
}
if (tid2eid.has_value()) {
STD_TORCH_CHECK(input_ids.has_value(),
"input_ids is required for hash MoE");
@@ -800,7 +751,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
tid2eid.value().const_data_ptr<int64_t>(), stream, is_padding_ptr);
tid2eid.value().const_data_ptr<int64_t>(), stream);
} else {
STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
torch::headeronly::ScalarType::Int);
@@ -810,7 +761,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int>(),
tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);
tid2eid.value().const_data_ptr<int>(), stream);
}
} else {
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
@@ -818,7 +769,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, false, static_cast<const OutIndType*>(nullptr),
static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);
static_cast<const OutIndType*>(nullptr), stream);
}
};
@@ -842,8 +793,7 @@ void topk_softplus_sqrt(
bool renormalize, double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& is_padding) {
const std::optional<torch::stable::Tensor>& tid2eid) {
const int num_experts = gating_output.size(-1);
const auto num_tokens = gating_output.numel() / num_experts;
const int topk = topk_weights.size(-1);
@@ -856,22 +806,21 @@ void topk_softplus_sqrt(
dispatch_topk_softplus_sqrt_launch<float>(
gating_output.const_data_ptr<float>(), topk_weights, topk_indices,
token_expert_indices, num_tokens, num_experts, topk, renormalize,
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream,
is_padding);
routed_scaling_factor, correction_bias, input_ids, tid2eid, stream);
} else if (gating_output.scalar_type() ==
torch::headeronly::ScalarType::Half) {
dispatch_topk_softplus_sqrt_launch<__half>(
reinterpret_cast<const __half*>(gating_output.const_data_ptr()),
topk_weights, topk_indices, token_expert_indices, num_tokens,
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, stream, is_padding);
input_ids, tid2eid, stream);
} else if (gating_output.scalar_type() ==
torch::headeronly::ScalarType::BFloat16) {
dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>(
reinterpret_cast<const __nv_bfloat16*>(gating_output.const_data_ptr()),
topk_weights, topk_indices, token_expert_indices, num_tokens,
num_experts, topk, renormalize, routed_scaling_factor, correction_bias,
input_ids, tid2eid, stream, is_padding);
input_ids, tid2eid, stream);
} else {
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
gating_output.scalar_type());
+3 -3
View File
@@ -8,19 +8,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
m.def(
"topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, Tensor? "
"bias, Tensor? is_padding) -> ()");
"bias) -> ()");
// Apply topk sigmoid to the gating outputs.
m.def(
"topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, "
"Tensor? bias, float routed_scaling_factor, Tensor? is_padding) -> ()");
"Tensor? bias, float routed_scaling_factor) -> ()");
m.def(
"topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! "
"token_expert_indices, Tensor gating_output, bool renormalize, float "
"routed_scaling_factor, Tensor? "
"bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()");
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
// Calculate the result of moe by summing up the partial results
// from all selected experts. topk_ids/expert_map are optional and, when
+12 -1
View File
@@ -294,6 +294,9 @@ FROM base AS rust-build
ARG BUILD_OS
ARG USE_SCCACHE
ARG SCCACHE_ENDPOINT
# Temporary default for the initial CI validation. Set this back to 0 when
# ci-infra passes VLLM_RUST_COVERAGE=1 explicitly.
ARG VLLM_RUST_COVERAGE=1
# Install native tools needed only for Rust/protoc builds.
RUN if [ "${BUILD_OS}" = "manylinux" ]; then \
@@ -793,7 +796,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
# Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
# https://docs.flashinfer.ai/installation.html
# From versions.json: .flashinfer.version
ARG FLASHINFER_VERSION=0.6.15.post1
ARG FLASHINFER_VERSION=0.6.14
RUN --mount=type=cache,target=/opt/uv/cache \
uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
--index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
@@ -902,6 +905,14 @@ COPY ./vllm/collect_env.py .
# note that this uses vllm installed by `pip`
FROM vllm-base AS test
COPY --from=rust-build \
/workspace/rust-coverage-tools/ \
/opt/vllm-rust-coverage/
ENV PATH=/opt/vllm-rust-coverage/bin:${PATH}
ENV LD_LIBRARY_PATH=/opt/vllm-rust-coverage/lib:${LD_LIBRARY_PATH}
ENV LLVM_PROFILE_FILE=/dev/null
ADD . /vllm-workspace/
ARG PYTHON_VERSION
+31 -53
View File
@@ -339,17 +339,18 @@ 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
# NIXL/UCX build stages
FROM base AS build_nixl
ARG NIXL_BRANCH="231d56753047c989062a5cb2ac703a1ad761c7d2"
ARG NIXL_REPO="https://github.com/ai-dynamo/nixl.git"
ARG UCX_BRANCH="96e58a16039f6d7d213bc967b8069238742c5194"
# RIXL/UCX build stages
FROM base AS build_rixl
ARG RIXL_BRANCH="39be1de8"
ARG RIXL_REPO="https://github.com/ROCm/RIXL.git"
ARG UCX_BRANCH="bfb51733"
ARG UCX_REPO="https://github.com/openucx/ucx.git"
ENV ROCM_PATH=/opt/rocm
ENV UCX_HOME=/usr/local/ucx
ENV NIXL_HOME=/usr/local/nixl
ENV RIXL_HOME=/usr/local/rixl
ENV RIXL_BENCH_HOME=/usr/local/rixl_bench
# NIXL build system dependencies and RDMA support
# RIXL build system dependences and RDMA support
RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
libgrpc-dev \
libgrpc++-dev \
@@ -367,8 +368,7 @@ RUN apt-get -y update && apt-get -y install autoconf libtool pkg-config \
&& rm -rf /var/lib/apt/lists/*
RUN --mount=type=cache,target=/root/.cache/uv \
uv pip install --system meson meson-python pybind11 pyyaml types-PyYAML \
auditwheel build patchelf pytest tomlkit "setuptools>=80.9.0"
uv pip install --system meson auditwheel patchelf tomlkit
RUN --mount=type=cache,target=/root/.cache/ccache \
cd /usr/local/src && \
@@ -396,50 +396,30 @@ ENV PATH=/usr/local/ucx/bin:$PATH
ENV LD_LIBRARY_PATH=${UCX_HOME}/lib:${LD_LIBRARY_PATH}
RUN --mount=type=cache,target=/root/.cache/ccache \
git clone ${NIXL_REPO} /opt/nixl && \
cd /opt/nixl && \
git checkout ${NIXL_BRANCH} && \
git clone ${RIXL_REPO} /opt/rixl && \
cd /opt/rixl && \
git checkout ${RIXL_BRANCH} && \
CC="ccache gcc" CXX="ccache g++" \
meson setup build --prefix=${NIXL_HOME} \
meson setup build --prefix=${RIXL_HOME} \
-Ducx_path=${UCX_HOME} \
-Dwheel_variant=rocm \
-Dbuild_tests=false \
-Dbuild_examples=false && \
-Drocm_path=${ROCM_PATH} && \
cd build && \
ninja -j$(nproc) && \
ninja install && \
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu" \
> /etc/ld.so.conf.d/nixl.conf && \
echo "${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins" \
>> /etc/ld.so.conf.d/nixl.conf && \
ldconfig
ninja install
# Generate the ROCm NIXL wheel. Upstream's generic wheel helper detects CUDA,
# so configure the ROCm wheel variant directly through Meson.
# Generate RIXL wheel
# Exclude libcore and libpull from auditwheel: transitive dependencies
# that are not shipped in the wheel and vary across base images.
RUN cd /opt/nixl && \
./contrib/tomlutil.py --wheel-name nixl-rocm pyproject.toml && \
CC="ccache gcc" CXX="ccache g++" \
uv build --wheel --no-build-isolation --out-dir /tmp/nixl_wheels \
--python ${PYTHON_VERSION} \
-Csetup-args=-Ducx_path=${UCX_HOME} \
-Csetup-args=-Dwheel_variant=rocm \
-Csetup-args=-Dbuild_tests=false \
-Csetup-args=-Dbuild_examples=false && \
mkdir -p /tmp/nixl_wheels/repaired /app/install && \
auditwheel repair \
--exclude 'libamdhip64*' \
--exclude 'libcore*' \
--exclude 'libpull*' \
/tmp/nixl_wheels/nixl_rocm*.whl \
--plat manylinux_2_34_$(uname -m) \
--wheel-dir /tmp/nixl_wheels/repaired && \
./contrib/wheel_add_ucx_plugins.py \
RUN cd /opt/rixl && \
sed -i "s/--exclude 'libamdhip64\*'/--exclude 'libamdhip64*' --exclude 'libcore*' --exclude 'libpull*'/" \
contrib/build-wheel.sh && \
mkdir -p /app/install && \
_ucx_install_dir=${UCX_HOME} \
./contrib/build-wheel.sh \
--output-dir /app/install \
--rocm-dir ${ROCM_PATH} \
--ucx-plugins-dir ${UCX_HOME}/lib/ucx \
--nixl-plugins-dir ${NIXL_HOME}/lib/$(uname -m)-linux-gnu/plugins \
/tmp/nixl_wheels/repaired/*.whl && \
cp /tmp/nixl_wheels/repaired/*.whl /app/install
--nixl-plugins-dir ${RIXL_HOME}/lib/x86_64-linux-gnu/plugins
# ROCShmem build stage - split from DeepEP so changing DEEPEP_BRANCH does not
# invalidate the slow ROCShmem build.
@@ -680,10 +660,10 @@ RUN if [ "${DEEPEP_NIC}" = "cx7" ] || [ "${DEEPEP_NIC}" = "io" ]; then \
ninja && ninja install && ldconfig && rm -rf /tmp/rdma-core; \
fi
# Install NIXL + DeepEP wheels.
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
# Install RIXL + DeepEP wheels.
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
--mount=type=bind,from=build_deepep,src=/app/deep_install,target=/deep_install \
uv pip install --system /nixl_install/*.whl /deep_install/*.whl
uv pip install --system /rixl_install/*.whl /deep_install/*.whl
# Copy ROCShmem runtime libraries.
COPY --from=build_rocshmem /opt/rocshmem /opt/rocshmem
@@ -744,7 +724,6 @@ ENV MIOPEN_DEBUG_CONV_GEMM=0
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc.
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
# ROCm profiler limits workaround.
RUN echo "ROCTRACER_MAX_EVENTS=10000000" > ${COMMON_WORKDIR}/libkineto.conf
@@ -817,9 +796,9 @@ RUN --mount=type=bind,from=export_vllm,src=/,target=/install \
&& pip uninstall -y vllm \
&& uv pip install --system *.whl
# Install NIXL ROCm wheel
RUN --mount=type=bind,from=build_nixl,src=/app/install,target=/nixl_install \
uv pip install --system /nixl_install/*.whl
# Install RIXL wheel
RUN --mount=type=bind,from=build_rixl,src=/app/install,target=/rixl_install \
uv pip install --system /rixl_install/*.whl
ARG COMMON_WORKDIR
ARG BASE_IMAGE
@@ -834,7 +813,6 @@ COPY --from=export_vllm /docker ${COMMON_WORKDIR}/vllm/docker
# Use legacy IPC mode for HSA to avoid GPU memory pinning issues with UCX rocm_ipc
# See: https://github.com/ROCm/rocm-libraries/issues/6266
ENV HSA_ENABLE_IPC_MODE_LEGACY=1
ENV UCX_RMA_PPLN_ENABLE=y
ENV TOKENIZERS_PARALLELISM=false
+1 -1
View File
@@ -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.post5"
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"
+1 -25
View File
@@ -86,29 +86,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
mkdir -p /tmp/hf-xet/dist && \
cp dist/*.whl /tmp/hf-xet/dist/
# Build LLVM 20 from source for llvmlite (system repos ship LLVM 21 which
# llvmlite v0.47 does not support; only SystemZ target is needed).
FROM base AS llvm20-build
ARG LLVM_VERSION=20.1.8
WORKDIR /tmp
RUN microdnf install -y ninja-build gcc gcc-c++ python3 xz && \
curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-${LLVM_VERSION}/llvm-project-${LLVM_VERSION}.src.tar.xz && \
tar -xf llvm-project-${LLVM_VERSION}.src.tar.xz && \
cmake -G Ninja -S llvm-project-${LLVM_VERSION}.src/llvm -B build \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_INSTALL_PREFIX=/opt/llvm20 \
-DLLVM_TARGETS_TO_BUILD="SystemZ" \
-DLLVM_ENABLE_RTTI=ON \
-DLLVM_BUILD_TOOLS=OFF \
-DLLVM_BUILD_UTILS=ON \
-DLLVM_BUILD_EXAMPLES=OFF \
-DLLVM_BUILD_TESTS=OFF \
-DLLVM_INCLUDE_TESTS=OFF \
-DLLVM_INCLUDE_EXAMPLES=OFF \
-DLLVM_INCLUDE_BENCHMARKS=OFF && \
ninja -C build install && \
rm -rf build llvm-project-${LLVM_VERSION}.src*
# Build numba
FROM python-install AS numba-builder
@@ -119,13 +96,11 @@ WORKDIR /tmp
# Clone all required dependencies
RUN --mount=type=cache,target=/root/.cache/uv \
--mount=type=bind,from=llvm20-build,source=/opt/llvm20,target=/opt/llvm20 \
microdnf install ninja-build gcc gcc-c++ -y && \
git clone --recursive https://github.com/numba/llvmlite.git -b v0.47.0 && \
git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
cd llvmlite && \
uv pip install 'cmake<4' 'setuptools<70' numpy && \
CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \
python setup.py bdist_wheel && \
cd ../numba && \
if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \
@@ -183,6 +158,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) && \
uv pip install -v \
$ARROW_WHL_FILE \
$VISION_WHL_FILE \
$HF_XET_WHL_FILE \
$LLVM_WHL_FILE \
+13 -13
View File
@@ -59,7 +59,7 @@ variable "PYTORCH_ROCM_ARCH" {
}
# Pre-built CI base image (Tier 1). Per-PR builds pull this instead of
# rebuilding NIXL/DeepEP/torchcodec from scratch. The ci_base stage in
# rebuilding RIXL/DeepEP/torchcodec from scratch. The ci_base stage in
# Dockerfile.rocm inherits from base, so CI_BASE_IMAGE only affects the test
# stage and is irrelevant when building --target ci_base itself.
variable "CI_BASE_IMAGE" {
@@ -75,7 +75,7 @@ variable "CI_MAX_JOBS" {
# Upstream dependency commit pins -- extracted from Dockerfile.rocm by
# ci-bake-rocm.sh at build time. Empty defaults are safe: the cache
# functions produce no entries when the variable is empty.
variable "NIXL_BRANCH" {
variable "RIXL_BRANCH" {
default = ""
}
@@ -91,7 +91,7 @@ variable "DEEPEP_BRANCH" {
default = ""
}
variable "NIXL_CACHE_KEY" {
variable "RIXL_CACHE_KEY" {
default = ""
}
@@ -236,7 +236,7 @@ function "get_cache_to_rocm_rust" {
])
}
# Cache functions for upstream dependency stages (NIXL/UCX, ROCShmem, DeepEP).
# Cache functions for upstream dependency stages (RIXL/UCX, ROCShmem, DeepEP).
# These stages are pinned to specific upstream commit hashes, so cache keys use
# those hashes rather than the Buildkite commit. This means the cache persists
# across all vLLM commits as long as the upstream dependency pins don't change.
@@ -244,16 +244,16 @@ function "get_cache_to_rocm_rust" {
function "get_cache_from_rocm_deps" {
params = []
result = compact([
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY}" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY}" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH}" : ""),
ROCSHMEM_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_CACHE_KEY}" : (ROCSHMEM_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rocshmem-rocm-${ROCSHMEM_BRANCH}" : ""),
DEEPEP_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_CACHE_KEY}" : (DEEPEP_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:deepep-rocm-${DEEPEP_BRANCH}-rocshmem-${ROCSHMEM_BRANCH}" : ""),
])
}
function "get_cache_to_rocm_nixl" {
function "get_cache_to_rocm_rixl" {
params = []
result = compact([
NIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_CACHE_KEY},mode=min" : (NIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:nixl-rocm-${NIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
RIXL_CACHE_KEY != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_CACHE_KEY},mode=min" : (RIXL_BRANCH != "" ? "type=registry,ref=${DOCKERHUB_CACHE_REPO}:rixl-rocm-${RIXL_BRANCH}-ucx-${UCX_BRANCH},mode=min" : ""),
])
}
@@ -372,11 +372,11 @@ variable "CI_BASE_IMAGE_TAG_STABLE" {
# in the registry cache keyed by its upstream commit hash. When ci_base rebuilds
# (e.g., requirements change), these stages are cache hits if their upstream
# pins haven't changed -- saving ~35min of compilation.
target "nixl-rocm-ci" {
target "rixl-rocm-ci" {
inherits = ["_common-rocm", "_ci-rocm"]
target = "build_nixl"
target = "build_rixl"
cache-from = get_cache_from_rocm_deps()
cache-to = get_cache_to_rocm_nixl()
cache-to = get_cache_to_rocm_rixl()
output = ["type=cacheonly"]
}
@@ -396,7 +396,7 @@ target "deepep-rocm-ci" {
output = ["type=cacheonly"]
}
# Builds only the ci_base stage (NIXL, DeepEP, torchcodec, etc.)
# Builds only the ci_base stage (RIXL, DeepEP, torchcodec, etc.)
# Invoked by the ensure-ci-base step when the content hash of ci_base-affecting
# files drifts from the remote image label. Per-PR builds then pull the result
# as CI_BASE_IMAGE instead of rebuilding those slow layers on every commit.
@@ -412,7 +412,7 @@ target "ci-base-rocm-ci" {
CI_BASE_IMAGE_TAG_CONTENT_EXTRA != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_CONTENT_EXTRA}" : "",
CI_BASE_IMAGE_TAG_STABLE != "" ? "type=registry,ref=${CI_BASE_IMAGE_TAG_STABLE}" : "",
]),
# Import upstream dependency caches so NIXL/ROCShmem/DeepEP stages
# Import upstream dependency caches so RIXL/ROCShmem/DeepEP stages
# are cache hits even when ci_base itself needs rebuilding.
get_cache_from_rocm_deps(),
)
@@ -424,5 +424,5 @@ target "ci-base-rocm-ci" {
# Group for ci_base builds -- exports dependency stage caches alongside the
# ci_base image so future rebuilds can reuse them independently.
group "ci-base-rocm-ci-with-deps" {
targets = ["nixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
targets = ["rixl-rocm-ci", "rocshmem-rocm-ci", "deepep-rocm-ci", "ci-base-rocm-ci"]
}
+2 -2
View File
@@ -53,7 +53,7 @@ variable "CI_BASE_IMAGE" {
# Upstream dependency commit pins. Plain local bake builds use the Dockerfile
# ARG defaults. ci-bake-rocm.sh resolves those defaults (plus any env
# overrides) and writes a small HCL override before invoking CI targets.
variable "NIXL_BRANCH" {
variable "RIXL_BRANCH" {
default = ""
}
@@ -106,7 +106,7 @@ target "test-rocm" {
output = ["type=docker"]
}
# CI base image target - builds only the ci_base stage (NIXL, DeepEP,
# CI base image target - builds only the ci_base stage (RIXL, DeepEP,
# torchcodec, requirements, etc.). Used by the weekly scheduled build and
# the auto-rebuild trigger when requirements change in a PR.
target "ci-base-rocm" {
+4 -1
View File
@@ -46,6 +46,9 @@
"TORCH_CUDA_ARCH_LIST": {
"default": "7.5 8.0 8.6 8.9 9.0 10.0 11.0 12.0"
},
"VLLM_RUST_COVERAGE": {
"default": "1"
},
"MAX_JOBS": {
"default": "2"
},
@@ -68,7 +71,7 @@
"default": "true"
},
"FLASHINFER_VERSION": {
"default": "0.6.15.post1"
"default": "0.6.14"
},
"GDRCOPY_CUDA_VERSION": {
"default": "12.8"
+5 -1
View File
@@ -13,7 +13,11 @@ Install the NIXL library: `uv pip install nixl`, as a quick start on Nvidia plat
- Refer to [NIXL official repository](https://github.com/ai-dynamo/nixl) for more installation instructions
- The specified required NIXL version can be found in [requirements/kv_connectors.txt](../../requirements/kv_connectors.txt) and other relevant config files
For ROCm, the [ROCm Dockerfile](../../docker/Dockerfile.rocm) builds NIXL and UCX with ROCm support from source.
For ROCm platform, the [ROCm docker file](../../docker/Dockerfile.rocm) includes RIXL and ucx already.
- Refer to [RIXL official repository](https://github.com/rocm/rixl) for more information
- The supportive libraries for RIXL can be found in [requirements/kv_connectors_rocm.txt](../../requirements/kv_connectors_rocm.txt)
- In the future we may remove RIXL from docker image file and users will be able to install from pre-compiled binary packages
For non-cuda platform, please install nixl with ucx build from source, instructed as below.
+1 -1
View File
@@ -315,7 +315,7 @@ vLLM CPU supports data parallel (DP), tensor parallel (TP) and pipeline parallel
- vLLM CPU supports quantizations:
- AWQ (x86 only)
- GPTQ (x86 only)
- compressed-tensor INT8 W8A8 (x86 only)
- compressed-tensor INT8 W8A8 (x86, s390x)
### Why do I see `get_mempolicy: Operation not permitted` when running in Docker?
@@ -11,7 +11,7 @@ Currently, the CPU implementation for s390x architecture supports FP32, BF16 and
- OS: `Linux`
- SDK: `gcc/g++ >= 14.0.0` or later with Command Line Tools
- Instruction Set Architecture (ISA): VXE support is required. Works with Z14 and above.
- Build from source python packages (no pre-built s390x wheels): `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`
- Build install python packages: `torchvision`, `llvmlite`, `numba`, `pyarrow (for testing)`, `opencv-headless`
--8<-- [end:requirements]
--8<-- [start:set-up-using-python]
@@ -28,24 +28,13 @@ Install the following packages from the package manager before building the vLLM
```bash
dnf install -y \
which procps findutils tar vim git patch xz ninja-build \
gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
which procps findutils tar vim git gcc-toolset-14 gcc-toolset-14-binutils gcc-toolset-14-libatomic-devel zlib-devel \
libjpeg-turbo-devel libtiff-devel libpng-devel libwebp-devel freetype-devel harfbuzz-devel \
openssl-devel openblas openblas-devel autoconf automake libtool cmake numpy libsndfile \
clang llvm-devel llvm-static clang-devel
```
Build and install `numactl` from source:
```bash
curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz
tar -xvzf v2.0.19.tar.gz
cd numactl-2.0.19
./autogen.sh && ./configure && make && make install
cd ..
```
Install rust>=1.80 which is needed for `outlines-core`, `uvloop`, and `hf-xet` python packages installation.
Install rust>=1.80 which is needed for `outlines-core` and `uvloop` python packages installation.
```bash
curl https://sh.rustup.rs -sSf | sh -s -- -y && \
@@ -55,79 +44,26 @@ curl https://sh.rustup.rs -sSf | sh -s -- -y && \
Execute the following commands to build and install vLLM from source.
!!! tip
Pre-built wheels are not available for s390x for the following packages. Build them from source before building vLLM: `torchvision`, `llvmlite`, `numba`, `opencv-python-headless`, `hf-xet`.
See `docker/Dockerfile.s390x` for exact versions and build commands used in each multi-stage build.
!!! note "LLVM 20 required for llvmlite"
`llvmlite v0.47` requires LLVM 20, but UBI 9.6 repos ship LLVM 21 which is
not compatible. You must build LLVM 20 from source before building `llvmlite`:
```bash
curl -LO https://github.com/llvm/llvm-project/releases/download/llvmorg-20.1.8/llvm-project-20.1.8.src.tar.xz
tar -xf llvm-project-20.1.8.src.tar.xz
cmake -G Ninja -S llvm-project-20.1.8.src/llvm -B llvm-build \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_INSTALL_PREFIX=/opt/llvm20 \
-DLLVM_TARGETS_TO_BUILD="SystemZ" \
-DLLVM_ENABLE_RTTI=ON \
-DLLVM_BUILD_TOOLS=OFF \
-DLLVM_BUILD_UTILS=ON \
-DLLVM_BUILD_EXAMPLES=OFF \
-DLLVM_BUILD_TESTS=OFF \
-DLLVM_INCLUDE_TESTS=OFF \
-DLLVM_INCLUDE_EXAMPLES=OFF \
-DLLVM_INCLUDE_BENCHMARKS=OFF
ninja -C llvm-build install
```
Then build `llvmlite` pointing to LLVM 20:
```bash
CMAKE_PREFIX_PATH=/opt/llvm20 LLVM_CONFIG=/opt/llvm20/bin/llvm-config \
python setup.py bdist_wheel
```
Please build the following dependencies, `torchvision`, `llvmlite`, `numba`, `llguidance`, `pyarrow`, `opencv-headless` from source before building vLLM.
```bash
uv pip install -v \
/path/to/torchvision.whl \
/path/to/llvmlite.whl \
/path/to/numba.whl \
/path/to/opencv_python_headless.whl \
/path/to/hf_xet.whl \
-r requirements/build/cpu.txt \
-r requirements/cpu.txt \
--torch-backend cpu \
--index-strategy unsafe-best-match && \
VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
uv pip install dist/*.whl
uv pip install -v \
-r requirements/build/cpu.txt \
-r requirements/cpu.txt \
--torch-backend cpu \
--index-strategy unsafe-best-match && \
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
uv pip install dist/*.whl
```
??? console "pip"
```bash
pip install -v \
--extra-index-url https://download.pytorch.org/whl/cpu \
/path/to/torchvision.whl \
/path/to/llvmlite.whl \
/path/to/numba.whl \
/path/to/opencv_python_headless.whl \
/path/to/hf_xet.whl \
-r requirements/build/cpu.txt \
-r requirements/cpu.txt && \
VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
pip install dist/*.whl
```
!!! warning "Protobuf workaround for s390x"
The C++ protobuf extension crashes on s390x. After installation, set the
following environment variable and remove the C++ extensions:
```bash
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
# Remove C++ protobuf extensions that crash on s390x
SITE_PKGS=$(python -c "import site; print(site.getsitepackages()[0])")
rm -rf "$SITE_PKGS/google/_upb/"*.so \
"$SITE_PKGS/google/protobuf/pyext/"*.so 2>/dev/null || true
pip install -v \
--extra-index-url https://download.pytorch.org/whl/cpu \
-r requirements/build/cpu.txt \
-r requirements/cpu.txt \
VLLM_TARGET_DEVICE=cpu python setup.py bdist_wheel && \
pip install dist/*.whl
```
--8<-- [end:build-wheel-from-source]
@@ -144,20 +80,19 @@ docker build -f docker/Dockerfile.s390x \
# Launch OpenAI server
docker run --rm \
--security-opt seccomp=unconfined \
--cap-add SYS_NICE \
--privileged true \
--shm-size 4g \
-p 8000:8000 \
-e VLLM_CPU_KVCACHE_SPACE=<KV cache space> \
-e VLLM_CPU_OMP_THREADS_BIND=<CPU cores for inference> \
vllm-cpu-env \
--model meta-llama/Llama-3.2-1B-Instruct \
--dtype bfloat16 \
--dtype float \
other vLLM OpenAI server arguments
```
!!! tip
Alternatively, `--privileged=true` also works but is broader and not generally recommended.
An alternative of `--privileged true` is `--cap-add SYS_NICE --security-opt seccomp=unconfined`.
--8<-- [end:build-image-from-source]
--8<-- [start:extra-information]
-3
View File
@@ -1,3 +0,0 @@
// Reo.Dev documentation tracking
// https://docs.reo.dev/integrations/input-sources/developer-insights/documentation
!function(){var e,t,n;e="d5c4337961ef0ac",t=function(){Reo.init({clientID:"d5c4337961ef0ac", enableThirdPartyTracking: true})},(n=document.createElement("script")).src="https://static.reo.dev/"+e+"/reo.js",n.defer=!0,n.onload=t,document.head.appendChild(n)}();
-1
View File
@@ -160,4 +160,3 @@ extra_javascript:
- https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
- mkdocs/javascript/edit_and_feedback.js
- mkdocs/javascript/slack_and_forum.js
- mkdocs/javascript/reo.js
+3 -3
View File
@@ -14,8 +14,8 @@ PyNvVideoCodec==2.0.4
# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
# install_requires so the published wheel does not carry an unresolvable pin
--extra-index-url https://flashinfer.ai/whl/
flashinfer-python==0.6.15.post1
flashinfer-cubin==0.6.15.post1
flashinfer-python==0.6.14
flashinfer-cubin==0.6.14
apache-tvm-ffi==0.1.10
tilelang==0.1.9
nvidia-cudnn-frontend>=1.19.1
@@ -26,7 +26,7 @@ fastsafetensors >= 0.3.2
# QuACK and Cutlass DSL for FA4 (cute-DSL implementation)
nvidia-cutlass-dsl[cu13]==4.6.0
quack-kernels>=0.6.1 # Required for CUTLASS DSL 4.6 by MSA
quack-kernels>=0.4.0 # Required for tml-fa4
# Tokenspeed_MLA for faster mla with spec decode
tokenspeed-mla==0.1.8; platform_system == "Linux"
+1 -1
View File
@@ -12,4 +12,4 @@ ray[data]
setuptools==78.1.0
setuptools-rust>=1.9.0
nixl==0.3.0
tpu-inference==0.25.0
tpu-inference==0.24.0
+2 -12
View File
@@ -5516,6 +5516,7 @@ dependencies = [
"asynk-strim-attr",
"bytes",
"clap",
"easy-ext",
"expect-test",
"futures",
"half",
@@ -5545,7 +5546,6 @@ dependencies = [
"tracing-subscriber",
"trait-set",
"uuid",
"vllm-chat-types",
"vllm-engine-core-client",
"vllm-llm",
"vllm-parser",
@@ -5555,16 +5555,6 @@ dependencies = [
"zeromq",
]
[[package]]
name = "vllm-chat-types"
version = "0.1.0"
dependencies = [
"easy-ext",
"serde",
"serde_json",
"serde_with",
]
[[package]]
name = "vllm-cmd"
version = "0.1.0"
@@ -5705,11 +5695,11 @@ dependencies = [
"expect-test",
"futures",
"openai-protocol",
"serde",
"serde_json",
"thiserror 2.0.18",
"thiserror-ext",
"tool-parser",
"vllm-chat-types",
"vllm-tokenizer",
"winnow",
"xgrammar-structural-tag",
-2
View File
@@ -2,7 +2,6 @@
members = [
"src/bench",
"src/chat",
"src/chat-types",
"src/cmd",
"src/engine-core-client",
"src/llm",
@@ -136,7 +135,6 @@ uuid = { version = "1.22.0", features = ["v4"] }
validator = { version = "0.20.0", features = ["derive"] }
vllm-bench = { path = "src/bench" }
vllm-chat = { path = "src/chat" }
vllm-chat-types = { path = "src/chat-types" }
vllm-engine-core-client = { path = "src/engine-core-client" }
vllm-llm = { path = "src/llm" }
vllm-managed-engine = { path = "src/managed-engine" }
-13
View File
@@ -14,10 +14,6 @@ service Generate {
rpc GenerateStream (GenerateRequest) returns (stream GenerateResponse) {}
}
service Control {
rpc Abort (AbortRequest) returns (AbortResponse) {}
}
// ======================================================================================
// Generate Request
// ======================================================================================
@@ -205,12 +201,3 @@ message TokenIds {
repeated uint32 ids = 1;
}
// ======================================================================================
// Control
// ======================================================================================
message AbortRequest {
repeated string request_ids = 1;
}
message AbortResponse {}
-14
View File
@@ -1,14 +0,0 @@
[package]
name = "vllm-chat-types"
version.workspace = true
edition.workspace = true
license.workspace = true
[dependencies]
easy-ext.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_with.workspace = true
[lints]
workspace = true
-155
View File
@@ -1,155 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::ops::Deref;
use serde::{Deserialize, Serialize};
/// One finalized assistant tool call.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AssistantToolCall {
/// Stable tool-call identifier.
pub id: String,
/// Function name selected by the assistant.
pub name: String,
/// Serialized function arguments.
pub arguments: String,
}
/// Semantic kind of one assistant output block.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AssistantBlockKind {
/// Visible final-answer text.
Text,
/// Extracted reasoning content.
Reasoning,
/// One finalized tool call.
ToolCall,
}
/// One structured assistant output block.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum AssistantContentBlock {
/// Visible final-answer text.
Text {
/// Visible text.
text: String,
},
/// Extracted reasoning content.
Reasoning {
/// Reasoning text.
text: String,
},
/// One finalized tool call.
ToolCall(AssistantToolCall),
}
impl AssistantContentBlock {
/// Return the semantic kind of this block.
pub fn kind(&self) -> AssistantBlockKind {
match self {
Self::Text { .. } => AssistantBlockKind::Text,
Self::Reasoning { .. } => AssistantBlockKind::Reasoning,
Self::ToolCall(..) => AssistantBlockKind::ToolCall,
}
}
/// Return this block as one finalized tool call when applicable.
pub fn as_tool_call(&self) -> Option<&AssistantToolCall> {
match self {
Self::ToolCall(call) => Some(call),
_ => None,
}
}
/// Trim whitespace from text and tool arguments.
///
/// Returns `None` when trimming makes a text or reasoning block empty.
pub fn trim(mut self) -> Option<Self> {
match &mut self {
Self::Text { text } | Self::Reasoning { text } => {
let trimmed_text = text.trim();
if trimmed_text.is_empty() {
return None;
}
*text = trimmed_text.to_string();
}
Self::ToolCall(call) => {
call.arguments = call.arguments.trim().to_string();
}
}
Some(self)
}
}
#[easy_ext::ext(AssistantMessageExt)]
impl [AssistantContentBlock] {
/// Concatenate all visible final-answer text blocks.
pub fn text(&self) -> String {
self.iter()
.filter_map(|block| match block {
AssistantContentBlock::Text { text } => Some(text.as_str()),
_ => None,
})
.collect()
}
/// Concatenate all extracted reasoning blocks.
pub fn reasoning(&self) -> Option<String> {
Some(
self.iter()
.filter_map(|block| match block {
AssistantContentBlock::Reasoning { text } => Some(text.as_str()),
_ => None,
})
.collect(),
)
.filter(|text: &String| !text.is_empty())
}
/// Return whether this assistant message contains reasoning text.
pub fn has_reasoning(&self) -> bool {
self.iter().any(|block| match block {
AssistantContentBlock::Reasoning { text } => !text.is_empty(),
_ => false,
})
}
/// Iterate over finalized assistant tool calls in encounter order.
pub fn tool_calls(&self) -> impl Iterator<Item = &AssistantToolCall> {
self.iter().filter_map(AssistantContentBlock::as_tool_call)
}
/// Return whether this assistant message contains any tool-call blocks.
pub fn has_tool_calls(&self) -> bool {
self.iter().any(|block| matches!(block, AssistantContentBlock::ToolCall(_)))
}
}
/// Final structured assistant message assembled from parsed output.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AssistantMessage {
/// Assistant content blocks in emission order.
pub content: Vec<AssistantContentBlock>,
}
impl Deref for AssistantMessage {
type Target = [AssistantContentBlock];
fn deref(&self) -> &Self::Target {
&self.content
}
}
impl AssistantMessage {
/// Push one new block to the end of the message content.
pub fn push_block(&mut self, block: AssistantContentBlock) {
self.content.push(block);
}
/// Trim all blocks and remove text blocks that become empty.
pub fn trim(mut self) -> Self {
self.content = self.content.into_iter().filter_map(AssistantContentBlock::trim).collect();
self
}
}
-191
View File
@@ -1,191 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use serde::{Deserialize, Serialize};
/// Detail level requested for an OpenAI-style image input.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ImageDetail {
/// Let the model-specific multimodal processor select the detail level.
#[default]
Auto,
/// Request low-detail image processing.
Low,
/// Request high-detail image processing.
High,
}
/// One chat content part in OpenAI-style block format.
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ChatContentPart {
/// One plain-text content block.
Text {
/// Plain-text content.
text: String,
},
/// One image URL or data URL content block.
ImageUrl {
/// Image URL or data URL.
image_url: String,
/// Requested image detail level.
detail: Option<ImageDetail>,
/// Optional caller-provided media identifier.
uuid: Option<String>,
},
/// One video URL or data URL content block.
VideoUrl {
/// Video URL or data URL.
video_url: String,
/// Optional caller-provided media identifier.
uuid: Option<String>,
},
/// One `input_audio` content block carrying base64-encoded audio bytes.
InputAudio {
/// Base64-encoded audio bytes.
data: String,
/// Optional audio format such as `wav` or `mp3`.
format: Option<String>,
/// Optional caller-provided media identifier.
uuid: Option<String>,
},
/// One audio URL or data URL content block.
AudioUrl {
/// Audio URL or data URL.
audio_url: String,
/// Optional caller-provided media identifier.
uuid: Option<String>,
},
}
impl ChatContentPart {
/// Construct one text content part with plain string content.
pub fn text(text: impl Into<String>) -> Self {
Self::Text { text: text.into() }
}
/// Construct one image URL content part with the given URL string.
pub fn image_url(image_url: impl Into<String>) -> Self {
Self::ImageUrl {
image_url: image_url.into(),
detail: None,
uuid: None,
}
}
/// Construct one video URL content part with the given URL string.
pub fn video_url(video_url: impl Into<String>) -> Self {
Self::VideoUrl {
video_url: video_url.into(),
uuid: None,
}
}
/// Construct one base64-encoded input-audio content part.
pub fn input_audio(data: impl Into<String>, format: Option<String>) -> Self {
Self::InputAudio {
data: data.into(),
format,
uuid: None,
}
}
/// Construct one audio URL content part with the given URL string.
pub fn audio_url(audio_url: impl Into<String>) -> Self {
Self::AudioUrl {
audio_url: audio_url.into(),
uuid: None,
}
}
/// Return the text content of this part.
///
/// Returns the static content-part type for multimodal content.
pub fn as_text(&self) -> Result<&str, &'static str> {
match self {
Self::Text { text } => Ok(text),
Self::ImageUrl { .. } => Err("image_url"),
Self::VideoUrl { .. } => Err("video_url"),
Self::InputAudio { .. } => Err("input_audio"),
Self::AudioUrl { .. } => Err("audio_url"),
}
}
/// Return whether this part is a text block with empty content.
fn is_empty_text(&self) -> bool {
matches!(self, Self::Text { text } if text.is_empty())
}
/// Return whether this part contains any multimodal content.
fn is_multimodal(&self) -> bool {
match self {
Self::Text { .. } => false,
Self::ImageUrl { .. }
| Self::VideoUrl { .. }
| Self::InputAudio { .. }
| Self::AudioUrl { .. } => true,
}
}
}
/// Chat content represented as a string or OpenAI-style content parts.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ChatContent {
/// Simple text content.
Text(String),
/// OpenAI-style content parts.
Parts(Vec<ChatContentPart>),
}
impl ChatContent {
/// Flatten text parts into one string without adding separators.
///
/// Returns the static content-part type when the content is multimodal.
pub fn try_flatten_to_text(&self) -> Result<String, &'static str> {
Ok(match self {
Self::Text(text) => text.clone(),
Self::Parts(parts) => parts
.iter()
.map(ChatContentPart::as_text)
.collect::<Result<Vec<_>, _>>()?
.concat(),
})
}
/// Return whether the content has no text or only empty text blocks.
pub fn is_empty(&self) -> bool {
match self {
Self::Text(text) => text.is_empty(),
Self::Parts(parts) => parts.iter().all(ChatContentPart::is_empty_text),
}
}
/// Return whether this content contains any multimodal parts.
pub fn has_multimodal(&self) -> bool {
match self {
Self::Text(_) => false,
Self::Parts(parts) => parts.iter().any(ChatContentPart::is_multimodal),
}
}
}
impl From<String> for ChatContent {
fn from(value: String) -> Self {
Self::Text(value)
}
}
impl From<&str> for ChatContent {
fn from(value: &str) -> Self {
Self::Text(value.to_string())
}
}
impl From<Vec<ChatContentPart>> for ChatContent {
fn from(value: Vec<ChatContentPart>) -> Self {
Self::Parts(value)
}
}
-26
View File
@@ -1,26 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
//! Engine-independent data types shared by chat renderers and output parsers.
//!
//! This crate defines chat history, rendering options, tool descriptions, and
//! structured assistant payloads. Serving requests, streamed events, renderer
//! implementations, parser state, and engine metadata live in their owning
//! crates.
mod assistant;
mod content;
mod message;
mod options;
#[cfg(test)]
mod tests;
mod tool;
pub use assistant::{
AssistantBlockKind, AssistantContentBlock, AssistantMessage, AssistantMessageExt,
AssistantToolCall,
};
pub use content::{ChatContent, ChatContentPart, ImageDetail};
pub use message::{ChatMessage, ChatRole};
pub use options::{ChatOptions, ChatToolChoice, GenerationPromptMode, ReasoningEffort};
pub use tool::Tool;
-195
View File
@@ -1,195 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use serde::{Deserialize, Serialize};
use crate::{AssistantContentBlock, AssistantMessage, AssistantMessageExt as _, ChatContent, Tool};
/// Role label for one chat message.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ChatRole {
/// System instructions.
System,
/// Developer instructions.
Developer,
/// User input.
User,
/// Assistant history.
Assistant,
/// Result of an assistant tool call.
ToolResponse,
}
impl ChatRole {
/// Return the role string exposed to chat templates.
pub fn as_str(&self) -> &'static str {
match self {
Self::System => "system",
Self::Developer => "developer",
Self::User => "user",
Self::Assistant => "assistant",
Self::ToolResponse => "tool_response",
}
}
}
/// One chat message.
///
/// Original Python API reference:
/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/entrypoints/chat_utils.py#L309-L333>
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "role", rename_all = "snake_case")]
pub enum ChatMessage {
/// System message.
System {
/// Message content.
content: ChatContent,
},
/// Developer message with optional message-local tools.
Developer {
/// Message content.
content: ChatContent,
/// Tools introduced by this developer message.
tools: Option<Vec<Tool>>,
},
/// User message.
User {
/// Message content.
content: ChatContent,
},
/// Assistant history assembled from structured blocks.
Assistant {
/// Structured assistant content.
content: Vec<AssistantContentBlock>,
},
/// Tool response associated with one prior assistant tool call.
ToolResponse {
/// Tool response content.
content: ChatContent,
/// Identifier of the assistant tool call being answered.
tool_call_id: String,
},
}
impl ChatMessage {
/// Construct one chat message with plain string content.
///
/// # Panics
///
/// Panics for [`ChatRole::ToolResponse`], which requires a tool-call ID.
/// Use [`Self::tool_response`] for tool responses.
pub fn text(role: ChatRole, text: impl Into<String>) -> Self {
let content: String = text.into();
match role {
ChatRole::System => Self::system(content),
ChatRole::Developer => Self::developer(content, None),
ChatRole::User => Self::user(content),
ChatRole::Assistant => Self::assistant_text(content),
ChatRole::ToolResponse => {
panic!(
"tool response messages require a tool_call_id; \
use ChatMessage::tool_response() instead"
)
}
}
}
/// Construct one system message.
pub fn system(content: impl Into<ChatContent>) -> Self {
Self::System {
content: content.into(),
}
}
/// Construct one developer message.
pub fn developer(content: impl Into<ChatContent>, tools: Option<Vec<Tool>>) -> Self {
Self::Developer {
content: content.into(),
tools,
}
}
/// Construct one user message.
pub fn user(content: impl Into<ChatContent>) -> Self {
Self::User {
content: content.into(),
}
}
/// Construct one assistant message with plain string content.
pub fn assistant_text(text: impl Into<String>) -> Self {
Self::Assistant {
content: vec![AssistantContentBlock::Text { text: text.into() }],
}
}
/// Construct one assistant message with structured content blocks.
pub fn assistant_blocks(content: Vec<AssistantContentBlock>) -> Self {
Self::Assistant { content }
}
/// Construct one tool-response message.
pub fn tool_response(content: impl Into<ChatContent>, tool_call_id: impl Into<String>) -> Self {
Self::ToolResponse {
content: content.into(),
tool_call_id: tool_call_id.into(),
}
}
/// Return the role of this message.
pub fn role(&self) -> ChatRole {
match self {
Self::System { .. } => ChatRole::System,
Self::Developer { .. } => ChatRole::Developer,
Self::User { .. } => ChatRole::User,
Self::Assistant { .. } => ChatRole::Assistant,
Self::ToolResponse { .. } => ChatRole::ToolResponse,
}
}
/// Concatenate the visible text carried by this message.
///
/// Returns the static content-part type when a non-assistant message
/// contains multimodal content.
pub fn text_content(&self) -> Result<String, &'static str> {
match self {
Self::System { content }
| Self::Developer { content, .. }
| Self::User { content }
| Self::ToolResponse { content, .. } => content.try_flatten_to_text(),
Self::Assistant { content } => Ok(content.text()),
}
}
/// Concatenate assistant reasoning text when present.
pub fn reasoning_content(&self) -> Option<String> {
match self {
Self::Assistant { content } => content.reasoning(),
Self::System { .. }
| Self::Developer { .. }
| Self::User { .. }
| Self::ToolResponse { .. } => None,
}
}
/// Return whether this message contains multimodal content.
pub fn has_multimodal(&self) -> bool {
match self {
Self::System { content }
| Self::Developer { content, .. }
| Self::User { content }
| Self::ToolResponse { content, .. } => content.has_multimodal(),
Self::Assistant { .. } => false,
}
}
}
impl From<AssistantMessage> for ChatMessage {
fn from(value: AssistantMessage) -> Self {
Self::Assistant {
content: value.content,
}
}
}
-134
View File
@@ -1,134 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
/// Controls how prompt rendering should end after the existing chat history.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum GenerationPromptMode {
/// Append a generation prompt for a new assistant turn.
///
/// Equivalent to `add_generation_prompt = true` and
/// `continue_final_message = false`.
#[default]
StartNewAssistant,
/// Leave the final assistant message open so generation continues it.
///
/// Equivalent to `add_generation_prompt = false` and
/// `continue_final_message = true`.
ContinueFinalAssistant,
/// Render the existing chat history without adding any trailing generation
/// prompt.
///
/// Equivalent to `add_generation_prompt = false` and
/// `continue_final_message = false`.
NoGenerationPrompt,
}
/// Effort level for reasoning models.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ReasoningEffort {
/// Disable reasoning.
None,
/// Use the smallest available reasoning effort.
Minimal,
/// Use low reasoning effort.
Low,
/// Use medium reasoning effort.
Medium,
/// Use high reasoning effort.
High,
/// Use extra-high reasoning effort.
XHigh,
/// Use the largest available reasoning effort.
Max,
}
impl ReasoningEffort {
/// Return the lowercase value exposed to chat templates.
pub fn as_str(self) -> &'static str {
match self {
Self::None => "none",
Self::Minimal => "minimal",
Self::Low => "low",
Self::Medium => "medium",
Self::High => "high",
Self::XHigh => "xhigh",
Self::Max => "max",
}
}
}
/// Chat-template-related request options.
///
/// These are the chat controls that currently affect prompt rendering.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ChatOptions {
/// Controls whether rendering starts a new assistant turn, continues the
/// final assistant message, or emits no trailing generation prompt.
pub generation_prompt_mode: GenerationPromptMode,
/// Per-request Jinja chat template override.
///
/// The renderer uses this template in place of the model's default chat
/// template when it is present.
pub chat_template: Option<String>,
/// Effort level exposed to chat templates for reasoning models.
pub reasoning_effort: Option<ReasoningEffort>,
/// Additional keyword arguments exposed to the chat template.
pub template_kwargs: HashMap<String, Value>,
}
impl Default for ChatOptions {
fn default() -> Self {
Self {
generation_prompt_mode: GenerationPromptMode::StartNewAssistant,
chat_template: None,
reasoning_effort: None,
template_kwargs: HashMap::new(),
}
}
}
impl ChatOptions {
/// Return whether rendering adds a prompt for a new assistant turn.
pub fn add_generation_prompt(&self) -> bool {
matches!(
self.generation_prompt_mode,
GenerationPromptMode::StartNewAssistant
)
}
/// Return whether rendering continues the final assistant message.
pub fn continue_final_message(&self) -> bool {
matches!(
self.generation_prompt_mode,
GenerationPromptMode::ContinueFinalAssistant
)
}
}
/// Tool-choice semantics supported by the shared chat types.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ChatToolChoice {
/// Disable tool calling.
#[default]
None,
/// Let the model choose whether to call a tool.
Auto,
/// Require the model to call a tool.
Required,
/// Require one named function.
Function {
/// Required function name.
name: String,
},
}
-117
View File
@@ -1,117 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use serde_json::{json, to_value};
use crate::{AssistantContentBlock, ChatContent, ChatContentPart, ChatMessage, ChatRole, Tool};
#[test]
fn chat_content_deserializes_from_raw_string() {
let content: ChatContent = serde_json::from_value(json!("hello")).unwrap();
assert_eq!(content, ChatContent::Text("hello".to_string()));
}
#[test]
fn chat_content_video_url_part_round_trips_through_serde() {
let content = ChatContent::Parts(vec![ChatContentPart::VideoUrl {
video_url: "https://example.com/demo.mp4".to_string(),
uuid: Some("video-1".to_string()),
}]);
let value = to_value(&content).unwrap();
assert_eq!(
value,
json!([{
"type": "video_url",
"video_url": "https://example.com/demo.mp4",
"uuid": "video-1",
}])
);
let decoded: ChatContent = serde_json::from_value(value).unwrap();
assert_eq!(decoded, content);
}
#[test]
fn chat_content_deserializes_from_openai_text_blocks() {
let content: ChatContent =
serde_json::from_value(json!([{ "type": "text", "text": "hello" }])).unwrap();
assert_eq!(
content,
ChatContent::Parts(vec![ChatContentPart::text("hello")])
);
}
#[test]
fn chat_content_from_string_like_values_builds_text() {
assert_eq!(
ChatContent::from("hello"),
ChatContent::Text("hello".to_string())
);
assert_eq!(
ChatContent::from("hello".to_string()),
ChatContent::Text("hello".to_string())
);
}
#[test]
fn chat_content_try_flattens_text_parts_without_separators() {
let content = ChatContent::Parts(vec![
ChatContentPart::text("hello"),
ChatContentPart::text(" world"),
]);
assert_eq!(content.try_flatten_to_text().unwrap(), "hello world");
}
#[test]
fn multimodal_content_parts_return_static_type_names() {
let parts = [
(ChatContentPart::image_url("image"), "image_url"),
(ChatContentPart::video_url("video"), "video_url"),
(ChatContentPart::input_audio("audio", None), "input_audio"),
(ChatContentPart::audio_url("audio"), "audio_url"),
];
for (part, expected) in parts {
assert_eq!(part.as_text(), Err(expected));
assert_eq!(
ChatContent::Parts(vec![part]).try_flatten_to_text(),
Err(expected)
);
}
}
#[test]
fn assistant_message_collects_visible_and_reasoning_text() {
let message = ChatMessage::assistant_blocks(vec![
AssistantContentBlock::Reasoning {
text: "inner".to_string(),
},
AssistantContentBlock::Text {
text: "outer".to_string(),
},
]);
assert_eq!(message.role(), ChatRole::Assistant);
assert_eq!(message.text_content().unwrap(), "outer");
assert_eq!(message.reasoning_content().as_deref(), Some("inner"));
}
#[test]
fn developer_message_round_trips_through_serde() {
let message = ChatMessage::developer(
"hello",
Some(vec![Tool {
name: "get_weather".to_string(),
description: Some("Get weather".to_string()),
parameters: json!({
"type": "object",
"properties": {"city": {"type": "string"}},
}),
strict: Some(true),
}]),
);
let value = to_value(&message).unwrap();
let decoded: ChatMessage = serde_json::from_value(value).unwrap();
assert_eq!(decoded, message);
}
-18
View File
@@ -1,18 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use serde::{Deserialize, Serialize};
use serde_json::Value;
/// One function-style tool made available to the model.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Tool {
/// Function name exposed to the model.
pub name: String,
/// Optional human-readable function description.
pub description: Option<String>,
/// JSON Schema describing the function parameters.
pub parameters: Value,
/// Optional strict-schema enforcement request.
pub strict: Option<bool>,
}
+1 -1
View File
@@ -7,6 +7,7 @@ license.workspace = true
[dependencies]
anyhow.workspace = true
asynk-strim-attr.workspace = true
easy-ext.workspace = true
futures.workspace = true
half.workspace = true
indexmap.workspace = true
@@ -29,7 +30,6 @@ tokio.workspace = true
tracing.workspace = true
trait-set.workspace = true
uuid.workspace = true
vllm-chat-types.workspace = true
vllm-engine-core-client.workspace = true
vllm-llm.workspace = true
vllm-parser.workspace = true
+146 -29
View File
@@ -1,17 +1,155 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::ops::Deref;
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use vllm_llm::TokenUsage;
use vllm_text::{DecodedLogprobs, DecodedPromptLogprobs};
use crate::FinishReason;
pub use vllm_chat_types::{
AssistantBlockKind, AssistantContentBlock, AssistantMessage, AssistantMessageExt,
AssistantToolCall,
};
/// One finalized assistant tool call.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AssistantToolCall {
pub id: String,
pub name: String,
pub arguments: String,
}
/// Semantic kind of one assistant output block.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AssistantBlockKind {
/// Visible final-answer text.
Text,
/// Extracted reasoning content.
Reasoning,
/// One finalized tool call.
ToolCall,
}
/// One structured assistant output block.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum AssistantContentBlock {
/// Visible final-answer text.
Text { text: String },
/// Extracted reasoning content.
Reasoning { text: String },
/// One finalized tool call.
ToolCall(AssistantToolCall),
}
impl AssistantContentBlock {
/// Return the semantic kind of this block.
pub fn kind(&self) -> AssistantBlockKind {
match self {
Self::Text { .. } => AssistantBlockKind::Text,
Self::Reasoning { .. } => AssistantBlockKind::Reasoning,
Self::ToolCall(..) => AssistantBlockKind::ToolCall,
}
}
/// Return this block as one finalized tool call, if applicable.
pub fn as_tool_call(&self) -> Option<&AssistantToolCall> {
match self {
Self::ToolCall(call) => Some(call),
_ => None,
}
}
/// Return a copy of this block with leading and trailing whitespace trimmed from all text
/// fields and tool call arguments, or `None` if the resulting text would be empty.
pub fn trim(mut self) -> Option<Self> {
match &mut self {
Self::Text { text } | Self::Reasoning { text } => {
let trimmed_text = text.trim();
if trimmed_text.is_empty() {
return None;
} else {
*text = trimmed_text.to_string();
}
}
Self::ToolCall(call) => {
call.arguments = call.arguments.trim().to_string();
}
}
Some(self)
}
}
#[easy_ext::ext(AssistantMessageExt)]
impl [AssistantContentBlock] {
/// Concatenate all visible final-answer text blocks.
pub fn text(&self) -> String {
self.iter()
.filter_map(|block| match block {
AssistantContentBlock::Text { text } => Some(text.as_str()),
_ => None,
})
.collect()
}
/// Concatenate all extracted reasoning blocks, if any.
pub fn reasoning(&self) -> Option<String> {
Some(
self.iter()
.filter_map(|block| match block {
AssistantContentBlock::Reasoning { text } => Some(text.as_str()),
_ => None,
})
.collect(),
)
.filter(|s: &String| !s.is_empty())
}
/// Return whether this assistant message contains any non-empty reasoning
/// text blocks.
pub fn has_reasoning(&self) -> bool {
self.iter().any(|block| match block {
AssistantContentBlock::Reasoning { text } => !text.is_empty(),
_ => false,
})
}
/// Return finalized assistant tool calls in encounter order.
pub fn tool_calls(&self) -> impl Iterator<Item = &AssistantToolCall> {
self.iter().filter_map(AssistantContentBlock::as_tool_call)
}
/// Return whether this assistant message contains any tool-call blocks.
pub fn has_tool_calls(&self) -> bool {
self.iter().any(|block| matches!(block, AssistantContentBlock::ToolCall(_)))
}
}
/// Final structured assistant message assembled from the event stream.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AssistantMessage {
pub content: Vec<AssistantContentBlock>,
}
impl Deref for AssistantMessage {
type Target = [AssistantContentBlock];
fn deref(&self) -> &Self::Target {
&self.content
}
}
impl AssistantMessage {
/// Push one new block to the end of the message content.
pub(crate) fn push_block(&mut self, block: AssistantContentBlock) {
self.content.push(block);
}
/// Return a copy of this message with leading and trailing whitespace trimmed from all text
/// fields and tool call arguments, and with any blocks that are empty after trimming removed.
pub fn trim(mut self) -> Self {
self.content = self.content.into_iter().filter_map(|block| block.trim()).collect();
self
}
}
/// Streamed chat event emitted by [`crate::ChatEventStream`].
#[derive(Debug, Clone, PartialEq)]
@@ -26,65 +164,44 @@ pub enum ChatEvent {
},
/// A new assistant output block has started.
BlockStart {
/// Stable block index within the assistant message.
index: usize,
/// Semantic kind of the opened block.
kind: AssistantBlockKind,
},
/// A newly observed delta for one open assistant output block.
BlockDelta {
/// Stable block index within the assistant message.
index: usize,
/// Semantic kind of the open block.
kind: AssistantBlockKind,
/// Newly emitted text.
delta: String,
},
/// Per-decoded-update sample metadata.
/// Per-decoded-update sample metadata: logprobs and/or output token IDs.
LogprobsDelta {
/// Decoded output logprobs, when requested.
logprobs: Option<DecodedLogprobs>,
/// Output token IDs emitted by this update.
token_ids: Vec<u32>,
},
/// One assistant output block has ended.
BlockEnd {
/// Stable block index within the assistant message.
index: usize,
/// Finalized block.
block: AssistantContentBlock,
},
/// One tool call has started.
ToolCallStart {
/// Stable tool-call index within the assistant message.
index: usize,
/// Stable tool-call identifier.
id: String,
/// Function name selected by the assistant.
name: String,
},
/// One incremental tool-call arguments delta.
ToolCallArgumentsDelta {
/// Stable tool-call index within the assistant message.
index: usize,
/// Newly emitted arguments text.
delta: String,
},
/// One incremental tool-call arguments delta for the currently open tool
/// call.
ToolCallArgumentsDelta { index: usize, delta: String },
/// One tool call has ended.
ToolCallEnd {
/// Stable tool-call index within the assistant message.
index: usize,
/// Finalized tool call.
call: AssistantToolCall,
},
/// Terminal event carrying the final assembled assistant message and finish
/// metadata.
Done {
/// Final structured assistant message.
message: AssistantMessage,
/// Final token usage.
usage: TokenUsage,
/// Reason generation stopped.
finish_reason: FinishReason,
/// Connector-specific KV transfer parameters for disaggregated serving.
kv_transfer_params: Option<serde_json::Value>,
+65 -110
View File
@@ -37,7 +37,7 @@ pub use renderer::{
};
pub use request::{
ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool,
ChatToolChoice, GenerationPromptMode, ImageDetail, ReasoningEffort, SamplingParams,
ChatToolChoice, GenerationPromptMode, ReasoningEffort, SamplingParams,
};
pub use stream::{ChatEventStream, ChatEventStreamTrait, CollectedAssistantMessage};
pub use vllm_llm::FinishReason;
@@ -54,7 +54,6 @@ mod stream;
use vllm_engine_core_client::EngineCoreClient;
use vllm_engine_core_client::protocol::dtype::ModelDtype;
use vllm_engine_core_client::protocol::multimodal::MmFeatures;
use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;
use vllm_llm::Llm;
use vllm_text::{Prompt, TextLlm, TextRequest};
@@ -89,92 +88,6 @@ pub fn validate_parser_overrides(
Ok(())
}
/// Chat request preparation shared by inference and render-only frontends.
pub struct ChatRequestProcessor {
backend: DynChatBackend,
/// Effective model dtype reported by the engine.
/// Absent for text-only frontends without an engine handshake.
model_dtype: Option<ModelDtype>,
}
impl ChatRequestProcessor {
/// Create a processor with multimodal support using the effective model
/// dtype reported by the engine.
fn new(backend: DynChatBackend, model_dtype: ModelDtype) -> Self {
Self {
backend,
model_dtype: Some(model_dtype),
}
}
/// Create a render-only processor that rejects multimodal requests.
pub fn render_only(backend: DynChatBackend) -> Self {
Self {
backend,
model_dtype: None,
}
}
async fn finalize_rendered_prompt(
&self,
request: &ChatRequest,
rendered: RenderedPrompt,
) -> Result<(Prompt, Option<MmFeatures>)> {
match self.model_dtype {
Some(model_dtype) => {
multimodal::finalize_rendered_prompt(
request,
rendered,
self.backend.multimodal_model_info(),
model_dtype,
)
.await
}
None if !request.has_multimodal() => Ok((rendered.prompt, None)),
None => Err(Error::UnsupportedMultimodalRenderer),
}
}
/// Prepare one chat request without submitting it to an engine.
pub async fn prepare(
&self,
mut request: ChatRequest,
options: NewChatOutputProcessorOptions<'_>,
) -> Result<(TextRequest, DynChatOutputProcessor)> {
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, options)?;
let rendered = self.backend.chat_renderer().render(&request)?;
let reasoning_parser_kwargs =
request
.sampling_params
.structured_outputs
.is_some()
.then(|| ReasoningParserKwargs {
chat_template_kwargs: rendered.effective_template_kwargs.clone(),
});
let (prompt, mm_features) = self.finalize_rendered_prompt(&request, rendered).await?;
let text_request = TextRequest {
request_id: request.request_id,
prompt,
mm_features,
sampling_params: request.sampling_params,
decode_options: request.decode_options,
intermediate: request.intermediate,
priority: request.priority,
cache_salt: request.cache_salt,
add_special_tokens: request.add_special_tokens,
data_parallel_rank: request.data_parallel_rank,
reasoning_parser_kwargs,
lora_request: request.lora_request,
arrival_time: Some(arrival_time),
};
Ok((text_request, output_processor))
}
}
/// Structured chat facade above [`TextLlm`].
///
/// This layer stays above raw text semantics: it takes care of chat-template
@@ -182,7 +95,9 @@ impl ChatRequestProcessor {
/// request semantics such as tool calls.
pub struct ChatLlm {
text: TextLlm,
processor: ChatRequestProcessor,
backend: DynChatBackend,
/// Effective model dtype reported by the engine.
model_dtype: ModelDtype,
/// Tool-call parser selection.
tool_call_parser: ParserSelection,
/// Reasoning parser selection.
@@ -197,7 +112,8 @@ impl ChatLlm {
Self {
text,
processor: ChatRequestProcessor::new(backend, model_dtype),
backend,
model_dtype,
tool_call_parser: ParserSelection::Auto,
reasoning_parser: ParserSelection::Auto,
}
@@ -224,7 +140,7 @@ impl ChatLlm {
/// Override the effective model dtype used for multimodal tensor encoding.
pub fn with_model_dtype(mut self, model_dtype: ModelDtype) -> Self {
self.processor.model_dtype = Some(model_dtype);
self.model_dtype = model_dtype;
self
}
@@ -256,36 +172,75 @@ impl ChatLlm {
}
/// Render, tokenize, and submit one chat request.
pub async fn chat(&self, request: ChatRequest) -> Result<ChatEventStream> {
let (text_request, output_processor) = self
.processor
.prepare(
request,
NewChatOutputProcessorOptions {
tool_call_parser: &self.tool_call_parser,
reasoning_parser: &self.reasoning_parser,
},
)
.await?;
let request_id = text_request.request_id.clone();
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 {
tool_call_parser: &self.tool_call_parser,
reasoning_parser: &self.reasoning_parser,
},
)?;
let rendered = self.backend.chat_renderer().render(&request)?;
let reasoning_parser_kwargs =
request
.sampling_params
.structured_outputs
.is_some()
.then(|| ReasoningParserKwargs {
chat_template_kwargs: rendered.effective_template_kwargs.clone(),
});
let (prompt, mm_features) = multimodal::finalize_rendered_prompt(
&request,
rendered,
self.backend.multimodal_model_info(),
self.model_dtype,
)
.await?;
let text_request = TextRequest {
request_id: request.request_id.clone(),
prompt,
mm_features,
sampling_params: request.sampling_params,
decode_options: request.decode_options,
intermediate: request.intermediate,
priority: request.priority,
cache_salt: request.cache_salt,
add_special_tokens: request.add_special_tokens,
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();
let structured_stream = output_processor.process(decoded_stream)?;
Ok(ChatEventStream::new(request_id, structured_stream))
Ok(ChatEventStream::new(request.request_id, structured_stream))
}
/// Render through the chat template and tokenize, without submitting to the engine.
///
/// Uses the same render, multimodal finalization, and encoding pipeline as
/// [`Self::chat`], but stops after token IDs so `/tokenize` counts match
/// what generation would see. Used by `POST /tokenize` (chat form).
/// Same render → [`multimodal::finalize_rendered_prompt`] → encode pipeline as
/// [`Self::chat`], but stops after token IDs so `/tokenize` counts match what
/// generation would see. Used by `POST /tokenize` (chat form).
pub async fn tokenize_chat(&self, request: ChatRequest) -> Result<Vec<u32>> {
request.validate()?;
let rendered = self.processor.backend.chat_renderer().render(&request)?;
let (prompt, _mm_features) =
self.processor.finalize_rendered_prompt(&request, rendered).await?;
let rendered = self.backend.chat_renderer().render(&request)?;
let (prompt, _mm_features) = multimodal::finalize_rendered_prompt(
&request,
rendered,
self.backend.multimodal_model_info(),
self.model_dtype,
)
.await?;
let tokenizer = self.text.tokenizer();
let token_ids = match prompt {
+2 -11
View File
@@ -33,7 +33,7 @@ use vllm_text::tokenizer::{DynTokenizer, Tokenizer};
use crate::error::{Error, Result, bail_multimodal, multimodal};
use crate::renderer::RenderedPrompt;
use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ImageDetail};
use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest};
mod audio;
mod expand;
@@ -529,7 +529,7 @@ fn extract_media_parts(request: &ChatRequest) -> Result<Vec<MediaContentPart>> {
uuid,
} => all_parts.push(MediaContentPart::ImageUrl {
url: image_url.clone(),
detail: detail.map(to_multimodal_image_detail),
detail: *detail,
uuid: uuid.clone(),
}),
ChatContentPart::VideoUrl { video_url, uuid } => {
@@ -556,15 +556,6 @@ fn extract_media_parts(request: &ChatRequest) -> Result<Vec<MediaContentPart>> {
Ok(all_parts)
}
/// Convert the protocol-level image detail into the multimodal processor type.
fn to_multimodal_image_detail(detail: ImageDetail) -> llm_multimodal::ImageDetail {
match detail {
ImageDetail::Auto => llm_multimodal::ImageDetail::Auto,
ImageDetail::Low => llm_multimodal::ImageDetail::Low,
ImageDetail::High => llm_multimodal::ImageDetail::High,
}
}
/// Wrap OpenAI base64 audio in a data URL consumed by `MediaConnector`.
fn input_audio_data_url(data: &str, format: Option<&str>) -> Result<String> {
let mime_type = match format {
+2 -2
View File
@@ -178,8 +178,8 @@ impl ChatOutputProcessor for DefaultChatOutputProcessor {
/// events through two sequential stages once text decoding has
/// already happened:
///
/// 1. `unified_event_stream` — reasoning and tool-call parsing
/// 2. `structured_chat_event_stream` — final block assembly
/// 1. [`unified_event_stream`] — reasoning and tool-call parsing
/// 2. [`structured_chat_event_stream`] — final block assembly
fn process(self: Box<Self>, decoded: DynDecodedTextEventStream) -> Result<DynChatEventStream> {
let parsed = unified_event_stream(decoded, self.parser);
let structured = structured_chat_event_stream(parsed, self.parallel_tool_calls);
@@ -517,7 +517,7 @@ fn write_chat_content(out: &mut String, content: &ChatContent) -> Result<()> {
ChatContent::Text(text) => out.push_str(text),
ChatContent::Parts(parts) => {
for part in parts {
out.push_str(part.as_text().map_err(Error::UnsupportedMultimodalContent)?);
out.push_str(part.as_text()?);
}
}
}
@@ -519,7 +519,7 @@ fn write_chat_content(out: &mut String, content: &ChatContent) -> Result<()> {
ChatContent::Text(text) => out.push_str(text),
ChatContent::Parts(parts) => {
for part in parts {
out.push_str(part.as_text().map_err(Error::UnsupportedMultimodalContent)?);
out.push_str(part.as_text()?);
}
}
}
+1 -1
View File
@@ -427,7 +427,7 @@ fn auto_drop_analysis_messages(messages: Vec<Message>) -> Vec<Message> {
/// Flatten vLLM text content and reject unsupported multimodal parts.
fn flatten_text(content: &ChatContent) -> Result<String> {
content.try_flatten_to_text().map_err(Error::UnsupportedMultimodalContent)
content.try_flatten_to_text()
}
/// Convert vLLM function tool definitions to Harmony tool descriptions.
+1 -1
View File
@@ -277,7 +277,7 @@ impl InklingChatRenderer {
tool_call_id: &str,
tool_call_id_to_name: &HashMap<String, String>,
) -> Result<()> {
let text = content.try_flatten_to_text().map_err(Error::UnsupportedMultimodalContent)?;
let text = content.try_flatten_to_text()?;
let tool_name = tool_call_id_to_name.get(tool_call_id).map(String::as_str).unwrap_or("");
self.write_text_block(out, self.special.message_tool, Some(tool_name), &text)
}
+546 -6
View File
@@ -1,17 +1,449 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::collections::HashMap;
use llm_multimodal::ImageDetail;
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub use vllm_chat_types::{
ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRole, ChatToolChoice,
GenerationPromptMode, ImageDetail, ReasoningEffort, Tool as ChatTool,
};
use vllm_engine_core_client::protocol::lora::LoraRequest;
pub use vllm_parser::tool::Tool as ChatTool;
pub use vllm_text::SamplingParams;
use vllm_text::TextDecodeOptions;
use crate::AssistantMessageExt;
use crate::error::{Error, Result};
use crate::event::{AssistantContentBlock, AssistantMessage};
/// Role label for one text-only chat message.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ChatRole {
System,
Developer,
User,
Assistant,
ToolResponse,
}
/// One text-only chat content part in OpenAI-style block format.
#[serde_with::skip_serializing_none]
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ChatContentPart {
/// One plain-text content block.
Text { text: String },
/// One image URL/data URL content block.
ImageUrl {
image_url: String,
detail: Option<ImageDetail>,
uuid: Option<String>,
},
/// One video URL/data URL content block.
VideoUrl {
video_url: String,
uuid: Option<String>,
},
/// One `input_audio` content block carrying base64-encoded audio bytes.
InputAudio {
data: String,
format: Option<String>,
uuid: Option<String>,
},
/// One audio URL/data URL content block.
AudioUrl {
audio_url: String,
uuid: Option<String>,
},
// ImageData...
// VideoData...
// ImageEmbeds...
}
impl ChatContentPart {
/// Construct one text content part with plain string content.
pub fn text(text: impl Into<String>) -> Self {
Self::Text { text: text.into() }
}
/// Construct one image URL content part with the given URL string.
pub fn image_url(image_url: impl Into<String>) -> Self {
Self::ImageUrl {
image_url: image_url.into(),
detail: None,
uuid: None,
}
}
/// Construct one video URL content part with the given URL string.
pub fn video_url(video_url: impl Into<String>) -> Self {
Self::VideoUrl {
video_url: video_url.into(),
uuid: None,
}
}
/// Construct one base64-encoded input-audio content part.
pub fn input_audio(data: impl Into<String>, format: Option<String>) -> Self {
Self::InputAudio {
data: data.into(),
format,
uuid: None,
}
}
/// Construct one audio URL content part with the given URL string.
pub fn audio_url(audio_url: impl Into<String>) -> Self {
Self::AudioUrl {
audio_url: audio_url.into(),
uuid: None,
}
}
/// Return the text content of this part when it's a text block, or an
/// "unsupported multimodal content" error otherwise.
pub(crate) fn as_text(&self) -> Result<&str> {
match self {
Self::Text { text } => Ok(text),
Self::ImageUrl { .. } => Err(Error::UnsupportedMultimodalContent("image_url")),
Self::VideoUrl { .. } => Err(Error::UnsupportedMultimodalContent("video_url")),
Self::InputAudio { .. } => Err(Error::UnsupportedMultimodalContent("input_audio")),
Self::AudioUrl { .. } => Err(Error::UnsupportedMultimodalContent("audio_url")),
}
}
/// Return whether this part is a text block with empty content.
pub(crate) fn is_empty_text(&self) -> bool {
matches!(self, Self::Text { text } if text.is_empty())
}
/// Return whether this part contains any multimodal content.
pub(crate) fn is_multimodal(&self) -> bool {
match self {
Self::Text { .. } => false,
Self::ImageUrl { .. }
| Self::VideoUrl { .. }
| Self::InputAudio { .. }
| Self::AudioUrl { .. } => true,
}
}
}
/// Text-only chat content.
///
/// This supports either a simple string or an OpenAI-style list of text blocks.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ChatContent {
/// Simple text content.
Text(String),
/// OpenAI-style blocks.
Parts(Vec<ChatContentPart>),
}
impl ChatContent {
/// Flatten the text content into one plain string without adding
/// separators.
// TODO: this method will be truly fallible once we add non-text content parts.
pub fn try_flatten_to_text(&self) -> Result<String> {
Ok(match self {
Self::Text(text) => text.clone(),
Self::Parts(parts) => {
parts.iter().map(ChatContentPart::as_text).collect::<Result<Vec<_>>>()?.concat()
}
})
}
/// Return whether there's no text content or only empty text blocks.
pub fn is_empty(&self) -> bool {
match self {
Self::Text(text) => text.is_empty(),
Self::Parts(parts) => parts.iter().all(ChatContentPart::is_empty_text),
}
}
/// Return whether this content contains any multimodal parts.
pub fn has_multimodal(&self) -> bool {
match self {
Self::Text(_) => false,
Self::Parts(parts) => parts.iter().any(ChatContentPart::is_multimodal),
}
}
}
impl From<String> for ChatContent {
fn from(value: String) -> Self {
Self::Text(value)
}
}
impl From<&str> for ChatContent {
fn from(value: &str) -> Self {
Self::Text(value.to_string())
}
}
impl From<Vec<ChatContentPart>> for ChatContent {
fn from(value: Vec<ChatContentPart>) -> Self {
Self::Parts(value)
}
}
/// One chat message.
///
/// Original Python API reference:
/// <https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/entrypoints/chat_utils.py#L309-L333>
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "role", rename_all = "snake_case")]
pub enum ChatMessage {
/// System message content.
System { content: ChatContent },
/// Developer message content plus optional message-local tools.
Developer {
content: ChatContent,
tools: Option<Vec<ChatTool>>,
},
/// User message content.
User { content: ChatContent },
/// Assistant history content assembled from structured assistant blocks.
Assistant { content: Vec<AssistantContentBlock> },
/// Tool response content associated with one prior assistant tool call.
ToolResponse {
content: ChatContent,
tool_call_id: String,
},
}
impl ChatMessage {
/// Construct one chat message with plain string content.
pub fn text(role: ChatRole, text: impl Into<String>) -> Self {
let content: String = text.into();
match role {
ChatRole::System => Self::system(content),
ChatRole::Developer => Self::developer(content, None),
ChatRole::User => Self::user(content),
ChatRole::Assistant => Self::assistant_text(content),
ChatRole::ToolResponse => {
panic!(
"tool response messages require a tool_call_id; \
use ChatMessage::tool_response() instead"
)
}
}
}
/// Construct one chat message with system role.
pub fn system(content: impl Into<ChatContent>) -> Self {
Self::System {
content: content.into(),
}
}
/// Construct one chat message with developer role.
pub fn developer(content: impl Into<ChatContent>, tools: Option<Vec<ChatTool>>) -> Self {
Self::Developer {
content: content.into(),
tools,
}
}
/// Construct one chat message with user role.
pub fn user(content: impl Into<ChatContent>) -> Self {
Self::User {
content: content.into(),
}
}
/// Construct one chat message with assistant role and plain string content.
pub fn assistant_text(text: impl Into<String>) -> Self {
Self::Assistant {
content: vec![AssistantContentBlock::Text { text: text.into() }],
}
}
/// Construct one chat message with assistant role and structured content
/// blocks.
pub fn assistant_blocks(content: Vec<AssistantContentBlock>) -> Self {
Self::Assistant { content }
}
/// Construct one tool-role message.
pub fn tool_response(content: impl Into<ChatContent>, tool_call_id: impl Into<String>) -> Self {
Self::ToolResponse {
content: content.into(),
tool_call_id: tool_call_id.into(),
}
}
/// Return the chat role of this message.
pub fn role(&self) -> ChatRole {
match self {
Self::System { .. } => ChatRole::System,
Self::Developer { .. } => ChatRole::Developer,
Self::User { .. } => ChatRole::User,
Self::Assistant { .. } => ChatRole::Assistant,
Self::ToolResponse { .. } => ChatRole::ToolResponse,
}
}
/// Concatenate the visible text carried by this message.
pub fn text_content(&self) -> Result<String> {
match self {
Self::System { content }
| Self::Developer { content, .. }
| Self::User { content }
| Self::ToolResponse { content, .. } => content.try_flatten_to_text(),
Self::Assistant { content } => Ok(content.text()),
}
}
/// Concatenate assistant reasoning text when present.
pub fn reasoning_content(&self) -> Option<String> {
match self {
Self::Assistant { content } => content.reasoning(),
Self::System { .. }
| Self::Developer { .. }
| Self::User { .. }
| Self::ToolResponse { .. } => None,
}
}
/// Return whether this message contains any multimodal content.
pub fn has_multimodal(&self) -> bool {
match self {
Self::System { content }
| Self::Developer { content, .. }
| Self::User { content }
| Self::ToolResponse { content, .. } => content.has_multimodal(),
Self::Assistant { .. } => false,
}
}
}
impl From<AssistantMessage> for ChatMessage {
fn from(value: AssistantMessage) -> Self {
Self::Assistant {
content: value.content,
}
}
}
/// Controls how prompt rendering should end after the existing chat history.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum GenerationPromptMode {
/// Append a generation prompt for a new assistant turn.
///
/// Equivalent to `add_generation_prompt = true` and `continue_final_message
/// = false`.
#[default]
StartNewAssistant,
/// Leave the final assistant message open so generation continues it.
///
/// Equivalent to `add_generation_prompt = false` and
/// `continue_final_message = true`.
ContinueFinalAssistant,
/// Render the existing chat history without adding any trailing generation
/// prompt.
///
/// Equivalent to `add_generation_prompt = false` and
/// `continue_final_message = false`.
NoGenerationPrompt,
}
/// Effort level for reasoning models.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ReasoningEffort {
None,
Minimal,
Low,
Medium,
High,
XHigh,
Max,
}
impl ReasoningEffort {
pub fn as_str(self) -> &'static str {
match self {
Self::None => "none",
Self::Minimal => "minimal",
Self::Low => "low",
Self::Medium => "medium",
Self::High => "high",
Self::XHigh => "xhigh",
Self::Max => "max",
}
}
}
/// Chat-template-related request options.
///
/// These are the small subset of chat controls that currently affect prompt
/// rendering in `vllm-chat`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ChatOptions {
/// Controls whether rendering starts a new assistant turn, continues the
/// final assistant message, or emits no trailing generation prompt at
/// all.
pub generation_prompt_mode: GenerationPromptMode,
/// Per-request Jinja chat template override. When set, this template is
/// used instead of the model's default chat template.
pub chat_template: Option<String>,
/// Effort level exposed to chat templates for reasoning models.
pub reasoning_effort: Option<ReasoningEffort>,
/// Additional keyword arguments exposed to the chat template.
pub template_kwargs: HashMap<String, Value>,
}
impl Default for ChatOptions {
fn default() -> Self {
Self {
generation_prompt_mode: GenerationPromptMode::StartNewAssistant,
chat_template: None,
reasoning_effort: None,
template_kwargs: HashMap::new(),
}
}
}
impl ChatOptions {
/// Whether to add a generation prompt for a new assistant turn after the
/// existing chat history.
pub fn add_generation_prompt(&self) -> bool {
matches!(
self.generation_prompt_mode,
GenerationPromptMode::StartNewAssistant
)
}
/// Whether to leave the final assistant message open so generation
/// continues it.
pub fn continue_final_message(&self) -> bool {
matches!(
self.generation_prompt_mode,
GenerationPromptMode::ContinueFinalAssistant
)
}
}
/// Tool-choice semantics supported by `vllm-chat`.
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ChatToolChoice {
#[default]
None,
Auto,
Required,
Function {
name: String,
},
}
/// One chat request ready to be rendered into a prompt and lowered into a
/// generate request.
@@ -146,12 +578,120 @@ impl ChatRequest {
}
}
impl ChatRole {
/// Return the chat-template role string used by the current text-only chat
/// backend.
pub fn as_str(&self) -> &'static str {
match self {
Self::System => "system",
Self::Developer => "developer",
Self::User => "user",
Self::Assistant => "assistant",
Self::ToolResponse => "tool_response",
}
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use serde_json::{json, to_value};
use super::ChatRequest;
use super::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatRole, ChatTool};
use crate::Error;
use crate::event::AssistantContentBlock;
#[test]
fn chat_content_deserializes_from_raw_string() {
let content: ChatContent = serde_json::from_value(json!("hello")).unwrap();
assert_eq!(content, ChatContent::Text("hello".to_string()));
}
#[test]
fn chat_content_video_url_part_round_trips_through_serde() {
let content = ChatContent::Parts(vec![ChatContentPart::VideoUrl {
video_url: "https://example.com/demo.mp4".to_string(),
uuid: Some("video-1".to_string()),
}]);
let value = to_value(&content).unwrap();
assert_eq!(
value,
json!([{
"type": "video_url",
"video_url": "https://example.com/demo.mp4",
"uuid": "video-1",
}])
);
let decoded: ChatContent = serde_json::from_value(value).unwrap();
assert_eq!(decoded, content);
}
#[test]
fn chat_content_deserializes_from_openai_text_blocks() {
let content: ChatContent =
serde_json::from_value(json!([{ "type": "text", "text": "hello" }])).unwrap();
assert_eq!(
content,
ChatContent::Parts(vec![ChatContentPart::text("hello")])
);
}
#[test]
fn chat_content_from_string_like_values_builds_text() {
assert_eq!(
ChatContent::from("hello"),
ChatContent::Text("hello".to_string())
);
assert_eq!(
ChatContent::from("hello".to_string()),
ChatContent::Text("hello".to_string())
);
}
#[test]
fn chat_content_try_flattens_text_parts_without_separators() {
let content = ChatContent::Parts(vec![
ChatContentPart::text("hello"),
ChatContentPart::text(" world"),
]);
assert_eq!(content.try_flatten_to_text().unwrap(), "hello world");
}
#[test]
fn assistant_message_collects_visible_and_reasoning_text() {
let message = ChatMessage::assistant_blocks(vec![
AssistantContentBlock::Reasoning {
text: "inner".to_string(),
},
AssistantContentBlock::Text {
text: "outer".to_string(),
},
]);
assert_eq!(message.role(), ChatRole::Assistant);
assert_eq!(message.text_content().unwrap(), "outer");
assert_eq!(message.reasoning_content().as_deref(), Some("inner"));
}
#[test]
fn developer_message_round_trips_through_serde() {
let message = ChatMessage::developer(
"hello",
Some(vec![ChatTool {
name: "get_weather".to_string(),
description: Some("Get weather".to_string()),
parameters: json!({
"type": "object",
"properties": {"city": {"type": "string"}},
}),
strict: Some(true),
}]),
);
let value = to_value(&message).unwrap();
let decoded: ChatMessage = serde_json::from_value(value).unwrap();
assert_eq!(decoded, message);
}
#[test]
fn enable_thinking_is_none_when_no_kwargs_are_present() {
+2 -2
View File
@@ -11,7 +11,7 @@ use tokio::time::timeout;
use vllm_chat::{
AssistantBlockKind, AssistantContentBlock, AssistantMessageExt as _, ChatBackend, ChatEvent,
ChatLlm, ChatMessage, ChatRenderer, ChatRequest, ChatRole, ChatTextBackend, ChatTool,
ChatToolChoice, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer, Error,
ChatToolChoice, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer,
FinishReason, GenerationPromptMode, NewChatOutputProcessorOptions, ParserSelection,
RenderedPrompt, SamplingParams,
};
@@ -257,7 +257,7 @@ impl ChatRenderer for FakeChatBackend {
for message in &request.messages {
prompt.push_str(message.role().as_str());
prompt.push_str(": ");
prompt.push_str(&message.text_content().map_err(Error::UnsupportedMultimodalContent)?);
prompt.push_str(&message.text_content()?);
prompt.push('\n');
}
if request.chat_options.add_generation_prompt() {
+1 -1
View File
@@ -9,10 +9,10 @@ test-util = []
[dependencies]
easy-ext.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
thiserror-ext.workspace = true
vllm-chat-types.workspace = true
vllm-tokenizer.workspace = true
winnow.workspace = true
xgrammar-structural-tag.workspace = true
+11 -1
View File
@@ -34,11 +34,21 @@ pub use minimax_m2::MinimaxM2ToolParser;
pub use minimax_m3::MinimaxM3ToolParser;
pub use qwen_coder::Qwen3CoderToolParser;
pub use seed_oss::SeedOssToolParser;
pub use vllm_chat_types::Tool;
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub use xgrammar_structural_tag::builders::StructuralTagBuilder;
use crate::utils;
/// One function-style tool made available to the model.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Tool {
pub name: String,
pub description: Option<String>,
pub parameters: Value,
pub strict: Option<bool>,
}
/// One tool-call update emitted while parsing assistant text.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolCallDelta {
+11 -6
View File
@@ -8,7 +8,7 @@ use tonic_health::ServingStatus;
use tonic_health::server::HealthReporter;
use tracing::{info, warn};
use super::{ControlGrpcService, GenerateGrpcService};
use super::GenerateGrpcService;
pub(crate) async fn monitor_health(
mut health_reporter: HealthReporter,
@@ -16,18 +16,21 @@ pub(crate) async fn monitor_health(
shutdown: CancellationToken,
) {
let generate_service = GenerateGrpcService::NAME;
let control_service = ControlGrpcService::NAME;
let status = ServingStatus::NotServing;
let health_event_first = tokio::select! {
result = engine_health.wait_for(|healthy| !*healthy) => {
match result {
Ok(_) => warn!(
generate_service,
overall_service = true,
status = ?status,
reason = "engine_unhealthy",
"marking gRPC health services as not serving"
),
Err(error) => warn!(
%error,
generate_service,
overall_service = true,
status = ?status,
reason = "health_channel_closed",
"engine health channel closed; marking gRPC health services as not serving"
@@ -37,6 +40,8 @@ pub(crate) async fn monitor_health(
}
_ = shutdown.cancelled() => {
info!(
generate_service,
overall_service = true,
status = ?status,
reason = "server_shutdown",
"server shutting down; marking gRPC health services as not serving"
@@ -46,20 +51,20 @@ pub(crate) async fn monitor_health(
};
health_reporter.set_not_serving::<GenerateGrpcService>().await;
health_reporter.set_not_serving::<ControlGrpcService>().await;
// Both gRPC services use the same engine client, so overall server health
// mirrors their shared engine health.
// Generate is currently the only engine-backed gRPC service, so overall
// server health intentionally mirrors it.
health_reporter.set_service_status("", status).await;
if health_event_first {
shutdown.cancelled().await;
info!(
generate_service,
overall_service = true,
reason = "server_shutdown",
"server shutting down; closing gRPC health watches"
);
}
health_reporter.clear_service_status(generate_service).await;
health_reporter.clear_service_status(control_service).await;
health_reporter.clear_service_status("").await;
}
-32
View File
@@ -26,10 +26,8 @@ pub mod pb {
}
pub(crate) use health::monitor_health;
pub use pb::control_server::ControlServer;
pub use pb::generate_server::GenerateServer;
pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>;
pub(crate) type GenerateGrpcService = GenerateServer<GenerateServiceImpl>;
#[cfg(test)]
@@ -46,36 +44,6 @@ impl GenerateServiceImpl {
}
}
/// gRPC control service backed by the shared application state.
pub struct ControlServiceImpl {
state: Arc<AppState>,
}
impl ControlServiceImpl {
pub fn new(state: Arc<AppState>) -> Self {
Self { state }
}
}
#[tonic::async_trait]
impl pb::control_server::Control for ControlServiceImpl {
async fn abort(
&self,
request: Request<pb::AbortRequest>,
) -> Result<Response<pb::AbortResponse>, Status> {
let request_ids = request.into_inner().request_ids;
if request_ids.is_empty() {
return Ok(Response::new(pb::AbortResponse {}));
}
self.state
.chat
.abort(&request_ids)
.await
.map_err(|error| Status::internal(error.to_report_string()))?;
Ok(Response::new(pb::AbortResponse {}))
}
}
#[tonic::async_trait]
impl pb::generate_server::Generate for GenerateServiceImpl {
type GenerateStreamStream =
+18 -122
View File
@@ -38,9 +38,8 @@ use vllm_tokenizer::test_utils::TestTokenizer;
use zeromq::prelude::{SocketRecv, SocketSend};
use zeromq::{DealerSocket, PushSocket, ZmqMessage};
use super::pb::control_client::ControlClient;
use super::pb::generate_client::GenerateClient;
use super::{ControlServer, ControlServiceImpl, GenerateServer, GenerateServiceImpl, pb};
use super::{GenerateServer, GenerateServiceImpl, pb};
use crate::listener::{Listener, MaybeTlsListener};
use crate::state::AppState;
use crate::tls;
@@ -154,6 +153,10 @@ async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec<bytes::Bytes> {
dealer.recv().await.expect("recv engine message").into_vec()
}
fn test_llm(client: EngineCoreClient) -> Llm {
Llm::new(client).with_request_id_randomization(false)
}
#[derive(Clone, Debug)]
struct FakeTextBackend;
@@ -203,7 +206,6 @@ async fn setup_grpc_service(
output_specs: Vec<(Vec<u32>, Option<EngineCoreFinishReason>)>,
) -> (
GenerateServer<GenerateServiceImpl>,
ControlServer<ControlServiceImpl>,
tokio::sync::watch::Receiver<bool>,
MockEngineTask,
) {
@@ -241,13 +243,12 @@ async fn setup_grpc_service(
let engine_health = client.subscribe_health();
let chat = ChatLlm::from_shared_backend(
Llm::new(client),
test_llm(client),
Arc::new(FakeTextBackend) as Arc<dyn ChatTextBackend>,
);
let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat));
(
GenerateServer::new(GenerateServiceImpl::new(state.clone())),
ControlServer::new(ControlServiceImpl::new(state)),
GenerateServer::new(GenerateServiceImpl::new(state)),
engine_health,
engine_task,
)
@@ -263,11 +264,9 @@ async fn grpc_test_server(
tokio::task::JoinHandle<()>,
MockEngineTask,
) {
let (generate_service, control_service, engine_health, engine_task) =
setup_grpc_service(engine_id, output_specs).await;
let (svc, engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let (channel, server_task) = start_grpc_test_server(
generate_service,
control_service,
svc,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
@@ -277,13 +276,11 @@ async fn grpc_test_server(
async fn start_grpc_test_server(
generate_service: GenerateServer<GenerateServiceImpl>,
control_service: ControlServer<ControlServiceImpl>,
engine_health: tokio::sync::watch::Receiver<bool>,
shutdown: tokio_util::sync::CancellationToken,
) -> (Channel, tokio::task::JoinHandle<()>) {
let (health_reporter, health_service) = health_reporter();
health_reporter.set_serving::<GenerateServer<GenerateServiceImpl>>().await;
health_reporter.set_serving::<ControlServer<ControlServiceImpl>>().await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
let addr = listener.local_addr().expect("local addr");
@@ -292,7 +289,6 @@ async fn start_grpc_test_server(
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
let server = TonicServer::builder()
.add_service(health_service)
.add_service(control_service)
.add_service(generate_service)
.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned());
let health_monitor =
@@ -323,8 +319,7 @@ async fn grpc_tls_test_server(
certs: &TestCerts,
cert_reqs: i32,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (generate_service, control_service, _engine_health, engine_task) =
setup_grpc_service(engine_id, output_specs).await;
let (svc, _engine_health, engine_task) = setup_grpc_service(engine_id, output_specs).await;
let context = tls::build_grpc_server_config(&server_tls(certs, cert_reqs))
.expect("build grpc tls config");
@@ -334,8 +329,7 @@ async fn grpc_tls_test_server(
let server_task = tokio::spawn(async move {
let incoming = MaybeTlsListener::tls(Listener::Tcp(listener), context);
TonicServer::builder()
.add_service(control_service)
.add_service(generate_service)
.add_service(svc)
.serve_with_incoming(incoming)
.await
.expect("grpc tls server");
@@ -415,7 +409,7 @@ async fn grpc_server_with_keepalive(
engine_id: impl Into<EngineId>,
keepalive: Option<Duration>,
) -> (String, tokio::task::JoinHandle<()>, MockEngineTask) {
let (generate_service, control_service, _engine_health, engine_task) =
let (svc, _engine_health, engine_task) =
setup_grpc_service(engine_id, default_stream_output_specs()).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.expect("bind grpc listener");
@@ -431,8 +425,7 @@ async fn grpc_server_with_keepalive(
let server_task = tokio::spawn(async move {
let incoming = MaybeTlsListener::plain(Listener::Tcp(listener));
builder
.add_service(control_service)
.add_service(generate_service)
.add_service(svc)
.serve_with_incoming(incoming)
.await
.expect("grpc server");
@@ -1080,106 +1073,14 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn control_abort_resolves_external_id_and_empty_is_noop() {
let (generate_service, control_service, engine_health, engine_task) =
setup_grpc_service(b"engine-grpc-abort-active", vec![(vec![b'h' as u32], None)]).await;
let (channel, server_task) = start_grpc_test_server(
generate_service,
control_service,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
.await;
let mut generate_client = GenerateClient::new(channel.clone());
let mut control_client = ControlClient::new(channel);
let request_id = "test-abort-active";
let mut stream = generate_client
.generate_stream(pb::GenerateRequest {
request_id: request_id.to_string(),
model: "test-model".to_string(),
prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
stopping: Some(pb::StoppingCriteria {
max_new_tokens: 10,
..Default::default()
}),
..Default::default()
})
.await
.expect("start generation")
.into_inner();
loop {
let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for active generation output")
.expect("read active generation output")
.expect("generation ended before producing output");
if let Some(output) = response.outputs {
assert!(
output.finish_info.is_none(),
"generation finished before abort behavior was exercised"
);
break;
}
}
control_client
.abort(pb::AbortRequest::default())
.await
.expect("empty abort should be a no-op");
assert!(
tokio::time::timeout(Duration::from_millis(100), stream.message())
.await
.is_err(),
"empty abort unexpectedly ended the active generation"
);
control_client
.abort(pb::AbortRequest {
request_ids: vec![
request_id.to_string(),
request_id.to_string(),
"unknown".to_string(),
],
})
.await
.expect("abort active generation");
let finish_reason = loop {
let response = tokio::time::timeout(Duration::from_secs(2), stream.message())
.await
.expect("timed out waiting for aborted generation")
.expect("read aborted generation")
.expect("generation ended without an aborted response");
if let Some(finish_info) = response.outputs.and_then(|output| output.finish_info) {
break finish_info.finish_reason;
}
};
assert_eq!(finish_reason, pb::finish_info::FinishReason::Aborted as i32);
control_client
.abort(pb::AbortRequest {
request_ids: vec![request_id.to_string()],
})
.await
.expect("repeated abort should be idempotent");
engine_task.await.expect("mock engine task");
server_task.abort();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy() {
let (generate_service, control_service, _connected_engine_health, _engine_task) =
let (generate_service, _connected_engine_health, _engine_task) =
setup_grpc_service(b"engine-grpc-health-failure", default_stream_output_specs()).await;
let (engine_health_tx, engine_health) = tokio::sync::watch::channel(true);
let (channel, server_task) = start_grpc_test_server(
generate_service,
control_service,
engine_health,
tokio_util::sync::CancellationToken::new(),
)
@@ -1187,7 +1088,7 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
let mut health_client = HealthClient::new(channel);
let mut health_streams = Vec::new();
for service in ["vllm.Generate", "vllm.Control", ""] {
for service in ["vllm.Generate", ""] {
let service_label = if service.is_empty() {
"overall"
} else {
@@ -1242,19 +1143,14 @@ async fn grpc_health_transitions_to_not_serving_when_engine_becomes_unhealthy()
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn grpc_health_watch_closes_on_graceful_shutdown() {
let (generate_service, control_service, engine_health, _engine_task) = setup_grpc_service(
let (generate_service, engine_health, _engine_task) = setup_grpc_service(
b"engine-grpc-health-shutdown",
default_stream_output_specs(),
)
.await;
let shutdown = tokio_util::sync::CancellationToken::new();
let (channel, server_task) = start_grpc_test_server(
generate_service,
control_service,
engine_health,
shutdown.clone(),
)
.await;
let (channel, server_task) =
start_grpc_test_server(generate_service, engine_health, shutdown.clone()).await;
let mut health_client = HealthClient::new(channel);
let mut stream = health_client
.watch(HealthCheckRequest {
-4
View File
@@ -207,9 +207,6 @@ where
let (health_reporter, health_service) = health_reporter();
let engine_health = state.engine_core_client().subscribe_health();
health_reporter.set_serving::<grpc::GenerateGrpcService>().await;
health_reporter.set_serving::<grpc::ControlGrpcService>().await;
let control_service =
grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone()));
let generate_service =
grpc::GenerateGrpcService::new(grpc::GenerateServiceImpl::new(state.clone()));
let svc = TonicServer::builder()
@@ -217,7 +214,6 @@ where
.http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT))
.layer(middleware::request_runtime_layer(state.clone()))
.add_service(health_service)
.add_service(control_service)
.add_service(generate_service);
info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server");
Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health))
@@ -19,7 +19,7 @@ use futures::StreamExt as _;
use serial_test::serial;
use vllm_chat::{
ChatBackend, ChatLlm, ChatRenderer, ChatRequest, ChatTextBackend, DefaultChatOutputProcessor,
DynChatOutputProcessor, DynChatRenderer, Error, NewChatOutputProcessorOptions, RenderedPrompt,
DynChatOutputProcessor, DynChatRenderer, NewChatOutputProcessorOptions, RenderedPrompt,
};
use vllm_engine_core_client::protocol::output::{
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs,
@@ -188,7 +188,7 @@ impl ChatRenderer for FakeChatBackend {
for message in &request.messages {
prompt.push_str(message.role().as_str());
prompt.push_str(": ");
prompt.push_str(&message.text_content().map_err(Error::UnsupportedMultimodalContent)?);
prompt.push_str(&message.text_content()?);
prompt.push('\n');
}
if request.chat_options.add_generation_prompt() {
@@ -399,10 +399,11 @@ mod tests {
use axum::http::HeaderMap;
use expect_test::expect;
use llm_multimodal::ImageDetail;
use serde_json::json;
use vllm_chat::{
AssistantContentBlock, AssistantToolCall, ChatContentPart, ChatMessage as VllmChatMessage,
ChatTool as VllmChatTool, ChatToolChoice, GenerationPromptMode, ImageDetail,
ChatTool as VllmChatTool, ChatToolChoice, GenerationPromptMode,
SamplingParams as VllmSamplingParams,
};
use vllm_text::output::TextDecodeOptions;
@@ -4,9 +4,9 @@
use std::collections::HashMap;
use std::slice;
use llm_multimodal::ImageDetail;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use vllm_chat::ImageDetail;
use vllm_llm::TokenUsage;
// ============================================================================
+2 -4
View File
@@ -23,7 +23,7 @@ use serial_test::serial;
use tower::{Service as _, ServiceExt as _};
use vllm_chat::{
ChatBackend, ChatContent, ChatContentPart, ChatLlm, ChatMessage, ChatRenderer, ChatRequest,
ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer, Error,
ChatTextBackend, DefaultChatOutputProcessor, DynChatOutputProcessor, DynChatRenderer,
NewChatOutputProcessorOptions,
};
use vllm_engine_core_client::mock_engine::default_ready_response;
@@ -539,9 +539,7 @@ fn render_fake_message_content(
| ChatMessage::Developer { content, .. }
| ChatMessage::User { content }
| ChatMessage::ToolResponse { content, .. } => render_fake_content(content, placeholder),
ChatMessage::Assistant { .. } => {
message.text_content().map_err(Error::UnsupportedMultimodalContent)
}
ChatMessage::Assistant { .. } => message.text_content(),
}
}
+51 -87
View File
@@ -38,22 +38,32 @@ trait_set! {
pub trait TextOutputStream = Stream<Item = Result<DecodedTextEvent>> + Send + 'static;
}
/// Text request preparation shared by inference and render-only frontends.
pub struct TextRequestProcessor {
/// Raw text facade above [`Llm`].
///
/// This layer stays below chat semantics: prompt text or prompt token IDs flow
/// in, decoded text deltas and terminal metadata flow out.
pub struct TextLlm {
/// Generate-only client owned by this text facade.
llm: Llm,
/// Tokenizer/model metadata backend responsible for prompt encode/decode
/// and sampling hints.
backend: DynTextBackend,
/// Runtime context window size reported by the engine startup handshake.
/// Render-only frontends supply the downstream engine's effective value.
max_model_len: u32,
/// Maximum number of top log probabilities accepted by this text facade.
max_logprobs: i32,
}
impl TextRequestProcessor {
/// Create a processor with the effective model context length.
pub fn new(backend: DynTextBackend, max_model_len: u32) -> Self {
impl TextLlm {
/// Create a new text-generation facade from a shared LLM client plus a text
/// backend.
pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
// The engine-reported value reflects the post-profiling, auto-fitted
// KV cache limit used at runtime.
let max_model_len = llm.engine_core_client().max_model_len();
Self {
llm,
backend,
max_model_len,
max_logprobs: SamplingLimits::DEFAULT_MAX_LOGPROBS,
@@ -68,83 +78,9 @@ impl TextRequestProcessor {
self
}
/// Return the tokenizer used by this processor.
pub fn tokenizer(&self) -> DynTokenizer {
self.backend.tokenizer()
}
/// Return the effective model context length.
pub fn max_model_len(&self) -> u32 {
self.max_model_len
}
/// Tokenize and lower one request without submitting it to an engine.
pub fn prepare(&self, mut request: TextRequest) -> Result<PreparedTextRequest> {
request.validate()?;
if request.arrival_time.is_none() {
request.arrival_time = Some(vllm_llm::current_unix_timestamp_secs());
}
let tokenizer = self.backend.tokenizer();
let prompt_token_ids = match take(&mut request.prompt) {
Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?,
// Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark
// and infra workloads bypass chat rendering and tokenizer overhead entirely.
Prompt::TokenIds(token_ids) => token_ids,
};
let sampling_hints = self.backend.sampling_hints()?;
let sampling_limits = SamplingLimits {
max_model_len: self.max_model_len,
max_logprobs: self.max_logprobs,
model_vocab_size: self.backend.model_vocab_size(),
tokenizer_vocab_size: self.backend.tokenizer_vocab_size(),
};
lower_text_request(
request,
prompt_token_ids,
sampling_hints,
sampling_limits,
tokenizer.as_ref(),
)
}
}
/// Raw text facade above [`Llm`].
///
/// This layer stays below chat semantics: prompt text or prompt token IDs flow
/// in, decoded text deltas and terminal metadata flow out.
pub struct TextLlm {
/// Generate-only client owned by this text facade.
llm: Llm,
/// Shared engine-free request preparation.
processor: TextRequestProcessor,
}
impl TextLlm {
/// Create a new text-generation facade from a shared LLM client plus a text
/// backend.
pub fn new(llm: Llm, backend: DynTextBackend) -> Self {
// The engine-reported value reflects the post-profiling, auto-fitted
// KV cache limit used at runtime.
let max_model_len = llm.engine_core_client().max_model_len();
Self {
llm,
processor: TextRequestProcessor::new(backend, max_model_len),
}
}
/// Override the maximum accepted logprobs count.
pub fn with_max_logprobs(mut self, max_logprobs: Option<i32>) -> Self {
self.processor = self.processor.with_max_logprobs(max_logprobs);
self
}
/// Return the backend model ID.
pub fn model_id(&self) -> &str {
self.processor.backend.model_id()
self.backend.model_id()
}
/// Expose the underlying engine-core client for low-level utility/admin
@@ -155,19 +91,19 @@ impl TextLlm {
/// Return the tokenizer used by this text backend.
pub fn tokenizer(&self) -> DynTokenizer {
self.processor.tokenizer()
self.backend.tokenizer()
}
/// Tokenizer vocabulary size (the number of tokens the tokenizer knows),
/// used to bound `allowed_token_ids` like the Python frontend `len(tokenizer)`.
pub fn tokenizer_vocab_size(&self) -> usize {
self.processor.backend.tokenizer_vocab_size()
self.backend.tokenizer_vocab_size()
}
/// Model vocabulary size from the model config, used to bound generated
/// token IDs and logits-domain sampling controls.
pub fn model_vocab_size(&self) -> usize {
self.processor.backend.model_vocab_size()
self.backend.model_vocab_size()
}
/// Tokenize if needed, lower to a generate request, and return the raw
@@ -181,7 +117,7 @@ impl TextLlm {
/// incrementally decoded text.
pub async fn generate(&self, request: TextRequest) -> Result<impl TextOutputStream> {
let (text_request, raw_stream) = self.generate_inner(request).await?;
let tokenizer = self.processor.tokenizer();
let tokenizer = self.backend.tokenizer();
let decoded_stream = output::decoded_text_event_stream(
text_request.request_id,
tokenizer,
@@ -195,12 +131,40 @@ impl TextLlm {
async fn generate_inner(
&self,
request: TextRequest,
mut request: TextRequest,
) -> Result<(TextRequest, GenerateOutputStream)> {
request.validate()?;
if request.arrival_time.is_none() {
request.arrival_time = Some(vllm_llm::current_unix_timestamp_secs());
}
let tokenizer = self.backend.tokenizer();
let prompt_token_ids = match take(&mut request.prompt) {
Prompt::Text(text) => tokenizer.encode(&text, request.add_special_tokens)?,
// Pre-tokenized prompts are the main completions-side escape hatch that lets benchmark
// and infra workloads bypass chat rendering and tokenizer overhead entirely.
Prompt::TokenIds(token_ids) => token_ids,
};
let sampling_hints = self.backend.sampling_hints()?;
let sampling_limits = SamplingLimits {
max_model_len: self.max_model_len,
max_logprobs: self.max_logprobs,
model_vocab_size: self.backend.model_vocab_size(),
tokenizer_vocab_size: self.backend.tokenizer_vocab_size(),
};
let PreparedTextRequest {
text_request,
generate_request,
} = self.processor.prepare(request)?;
} = lower_text_request(
request,
prompt_token_ids,
sampling_hints,
sampling_limits,
&*tokenizer,
)?;
let raw_stream = self.llm.generate(generate_request).await?;
Ok((text_request, raw_stream))
-62
View File
@@ -214,65 +214,3 @@ def test_cache_config_hash_ignores_kv_cache_sizing_knobs():
base_hash = CacheConfig().compute_hash()
assert CacheConfig(kv_cache_memory_bytes=1 << 30).compute_hash() == base_hash
assert CacheConfig(gpu_memory_utilization=0.5).compute_hash() == base_hash
def test_envs_compile_factors_relocation_invariant(tmp_path):
"""Relocating HOME or the XDG roots must not change the compile-cache
env hash.
Location-derived env vars (VLLM_XLA_CACHE_PATH from XDG_CACHE_HOME,
VLLM_CONFIG_ROOT from XDG_CONFIG_HOME/HOME) carry no information about
compiled artifacts, only about where directories live. When they leak
into compile_factors(), a cache produced under one HOME/XDG layout
silently misses under another - which defeats copying or pre-baking a
compile cache into a container image.
"""
import os
import subprocess
import sys
code = """
import sys
import logging
logging.disable(logging.CRITICAL)
from vllm import envs
from vllm.config.utils import hash_factors
print(hash_factors(envs.compile_factors()))
"""
def hash_with(extra_env):
env = {**dict(os.environ), "VLLM_LOGGING_LEVEL": "ERROR"}
# Drop explicit overrides so the derived defaults are what is
# exercised, then apply the relocation under test.
for key in ("VLLM_XLA_CACHE_PATH", "VLLM_CONFIG_ROOT", "VLLM_CACHE_ROOT"):
env.pop(key, None)
env.update(extra_env)
result = subprocess.run(
[sys.executable, "-c", code],
capture_output=True,
text=True,
check=True,
env=env,
)
return result.stdout.strip()
xdg_cache = tmp_path / "relocated-xdg-cache"
xdg_config = tmp_path / "relocated-xdg-config"
new_home = tmp_path / "relocated-home"
for d in (xdg_cache, xdg_config, new_home):
d.mkdir()
base = hash_with({})
relocated_xdg = hash_with(
{"XDG_CACHE_HOME": str(xdg_cache), "XDG_CONFIG_HOME": str(xdg_config)}
)
relocated_home = hash_with({"HOME": str(new_home)})
assert relocated_xdg == base, (
"XDG_CACHE_HOME/XDG_CONFIG_HOME relocation changed the compile-cache "
"env hash - a location-only derived var is leaking into the key"
)
assert relocated_home == base, (
"HOME relocation changed the compile-cache env hash - a "
"location-only derived var is leaking into the key"
)
-9
View File
@@ -43,15 +43,6 @@ def test_language_model_only_affects_model_hash():
assert base_hash != lm_only_hash
@pytest.mark.parametrize("backend_arg", ["video_backend", "backend"])
def test_use_gpu_video_backend_from_media_io_kwargs(backend_arg: str):
config = MultiModalConfig(
media_io_kwargs={"video": {backend_arg: "pynvvideocodec"}}
)
assert config.use_gpu_video_backend()
def test_mm_encoder_fp8_scale_path_requires_fp8():
with pytest.raises(ValueError, match="mm_encoder_attn_dtype"):
MultiModalConfig(mm_encoder_fp8_scale_path="/tmp/scales.json")
+5 -8
View File
@@ -254,8 +254,8 @@ def test_multiproc_executor_shutdown_cleanup():
for worker in executor.workers:
assert not worker.proc.is_alive(), "Worker processes should be terminated"
# Verify shutdown flag is set
assert executor.shutting_down, "Shutdown flag should be set"
# Verify shutdown event is set
assert executor.shutdown_event.is_set(), "Shutdown event should be set"
# Multiple shutdowns should be safe (idempotent)
executor.shutdown()
@@ -292,6 +292,7 @@ def test_multiproc_executor_pipeline_parallel():
"Max concurrent batches should follow the configured PP/async "
"scheduling policy"
)
finally:
# Clean up
executor.shutdown()
@@ -337,10 +338,6 @@ def test_multiproc_executor_multi_node():
- Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2
Total world_size = 4, nnodes = 2
"""
# Python 3.14+ changed default multiprocessing start method to 'forkserver'
# which cannot pickle nested functions. Use 'fork' for this test.
mp_ctx = multiprocessing.get_context("fork")
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
port = s.getsockname()[1]
@@ -408,12 +405,12 @@ def test_multiproc_executor_multi_node():
executor.shutdown()
# Create a queue to collect results from both processes
result_queue: multiprocessing.Queue[dict[str, int | bool]] = mp_ctx.Queue()
result_queue: multiprocessing.Queue[dict[str, int | bool]] = multiprocessing.Queue()
# Start both node processes
processes = []
for node_rank in range(2):
p = mp_ctx.Process(
p = multiprocessing.Process(
target=run_node,
args=(node_rank, result_queue, port),
name=f"Node{node_rank}",
@@ -9,6 +9,8 @@ from pydantic import TypeAdapter, ValidationError
from vllm import PoolingParams
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
from vllm.entrypoints.pooling.embed.protocol import (
CohereEmbedContent,
CohereEmbedInput,
CohereEmbedRequest,
EmbeddingBatchChatInputRequest,
EmbeddingBatchChatRequest,
@@ -17,10 +19,7 @@ from vllm.entrypoints.pooling.embed.protocol import (
EmbeddingCompletionRequest,
EmbeddingRequest,
)
from vllm.entrypoints.pooling.typing import (
PoolingEngineInput,
PoolingServeContext,
)
from vllm.entrypoints.pooling.typing import PoolingServeContext
from vllm.outputs import PoolingOutput, PoolingRequestOutput
@@ -411,7 +410,6 @@ class TestChunkedEmbeddingProcessing:
def _make_handler(cls):
handler = object.__new__(EmbedIOProcessor)
handler.model_config = cls._FakeModelConfig()
handler.enable_chunked_processing = True
return handler
@staticmethod
@@ -423,29 +421,15 @@ class TestChunkedEmbeddingProcessing:
}
)
assert isinstance(request, EmbeddingCompletionRequest)
pooling_params = PoolingParams()
return PoolingServeContext(
request=request,
pooling_params=pooling_params,
pooling_params=PoolingParams(),
model_name="test",
request_id="embd-client-prompt-999-chunk-888",
engine_inputs=[
PoolingEngineInput(
prompts={"prompt_token_ids": [0, 1, 2, 3, 4]},
params=pooling_params,
lora_requests=None,
priorities=0,
),
PoolingEngineInput(
prompts={"prompt_token_ids": [10, 11]},
params=pooling_params,
lora_requests=None,
priorities=0,
),
{"prompt_token_ids": [0, 1, 2, 3, 4]},
{"prompt_token_ids": [10, 11]},
],
lora_request=None,
priorities=0,
prompt_extras=None,
)
@staticmethod
@@ -466,7 +450,7 @@ class TestChunkedEmbeddingProcessing:
handler = self._make_handler()
ctx = self._make_context()
handler.maybe_pre_process_chunked(ctx)
handler._pre_process_chunked(ctx)
assert ctx.prompt_request_ids == [
"embd-client-prompt-999-chunk-888-prompt-0-chunk-0",
@@ -504,3 +488,227 @@ class TestChunkedEmbeddingProcessing:
ctx.final_res_batch[1].outputs.data,
torch.tensor([9.0, 9.0]),
)
class TestPreProcessCohereOnline:
"""Unit tests for EmbedIOProcessor._pre_process_cohere_online."""
@staticmethod
def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]:
return PoolingServeContext(
request=CohereEmbedRequest(model="test", **request_kwargs),
pooling_params=PoolingParams(),
model_name="test",
request_id="embd-test",
)
@staticmethod
def _make_handler():
handler = object.__new__(EmbedIOProcessor)
handler._validate_input_type = lambda _input_type: None
return handler
def test_text_only_without_task_prefix_uses_completion_path(self):
handler = self._make_handler()
ctx = self._make_context(texts=["hello"])
calls: list[tuple[str, object]] = []
def preprocess_cmpl_online(request, prompt_input, prompt_embeds):
calls.append(("completion", prompt_input))
return ["completion"]
handler._get_task_instruction_prefix = lambda _input_type: None
handler._has_chat_template = lambda: False
handler._preprocess_cmpl_online = preprocess_cmpl_online
handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail(
"text-only request should not require chat rendering"
)
handler._pre_process_cohere_online(ctx)
assert ctx.engine_inputs == ["completion"]
assert calls == [("completion", ["hello"])]
def test_text_only_falls_back_to_prefixed_completion_without_template(self):
handler = self._make_handler()
ctx = self._make_context(texts=["hello"], input_type="query")
calls: list[tuple[str, object]] = []
def preprocess_cmpl(request, prompt_input, prompt_embeds):
calls.append(("completion", prompt_input))
return ["fallback"]
handler._get_task_instruction_prefix = lambda _input_type: "query: "
handler._has_chat_template = lambda: False
handler._batch_render_chat = lambda *_args, **_kwargs: pytest.fail(
"chat rendering should be skipped without a template"
)
handler._preprocess_cmpl_online = preprocess_cmpl
handler._pre_process_cohere_online(ctx)
assert ctx.engine_inputs == ["fallback"]
assert calls == [("completion", ["query: hello"])]
def test_text_only_with_template_uses_chat_path(self):
handler = self._make_handler()
ctx = self._make_context(texts=["hello"], input_type="query")
calls: list[tuple[str, object]] = []
def batch_render_chat(
request,
all_messages,
truncate_prompt_tokens,
truncation_side,
):
calls.append(
(
"chat",
{
"request": request,
"all_messages": all_messages,
"truncate_prompt_tokens": truncate_prompt_tokens,
"truncation_side": truncation_side,
},
)
)
return ["chat"]
handler._get_task_instruction_prefix = lambda _input_type: "query: "
handler._has_chat_template = lambda: True
handler._batch_render_chat = batch_render_chat
handler._preprocess_cmpl_online = lambda *_args, **_kwargs: pytest.fail(
"completion path should be skipped when a template exists"
)
handler._pre_process_cohere_online(ctx)
assert ctx.engine_inputs == ["chat"]
assert calls == [
(
"chat",
{
"request": ctx.request,
"all_messages": [
handler._mixed_input_to_messages(
CohereEmbedInput(
content=[CohereEmbedContent(type="text", text="hello")]
),
task_prefix="query: ",
)
],
"truncate_prompt_tokens": -1,
"truncation_side": None,
},
)
]
class TestPreProcessOpenAIEmbeddingChatOnline:
"""Unit tests for OpenAI embedding chat preprocessing."""
class _FakeModelConfig:
max_model_len = 128
encoder_config: dict[str, object] = {}
pooler_config = None
multimodal_config = None
is_encoder_decoder = False
class _FakeRenderer:
tokenizer = object()
def __init__(self):
self.calls = []
def render_chat(
self,
all_messages,
chat_params,
tok_params,
prompt_extras=None,
):
self.calls.append(
{
"all_messages": all_messages,
"chat_params": chat_params,
"tok_params": tok_params,
"prompt_extras": prompt_extras,
}
)
return all_messages, [
{"prompt_token_ids": [index]} for index, _ in enumerate(all_messages)
]
@classmethod
def _make_handler(cls, renderer):
handler = object.__new__(EmbedIOProcessor)
handler.renderer = renderer
handler.model_config = cls._FakeModelConfig()
handler.chat_template = "template"
handler.chat_template_content_format = "auto"
handler.trust_request_chat_template = False
handler.enable_chunked_processing = False
return handler
@staticmethod
def _make_context(
request: (
EmbeddingChatRequest
| EmbeddingBatchChatRequest
| EmbeddingChatInputRequest
| EmbeddingBatchChatInputRequest
),
) -> PoolingServeContext[
EmbeddingChatRequest
| EmbeddingBatchChatRequest
| EmbeddingChatInputRequest
| EmbeddingBatchChatInputRequest
]:
return PoolingServeContext(
request=request,
pooling_params=PoolingParams(),
model_name="test",
request_id="embd-test",
)
def test_chat_template_kwargs_forwarded_for_batched_input_messages(self):
request = TypeAdapter(EmbeddingRequest).validate_python(
{
"model": "test",
"input": [
[{"role": "user", "content": "hello"}],
[{"role": "user", "content": "goodbye"}],
],
"add_generation_prompt": True,
"chat_template_kwargs": {"instruction": "Represent the query: "},
"mm_processor_kwargs": {"max_pixels": 1},
"cache_salt": "salt",
}
)
assert isinstance(request, EmbeddingBatchChatInputRequest)
renderer = self._FakeRenderer()
handler = self._make_handler(renderer)
ctx = self._make_context(request)
handler.pre_process_online(ctx)
assert ctx.engine_inputs == [
{"prompt_token_ids": [0]},
{"prompt_token_ids": [1]},
]
assert len(renderer.calls) == 1
call = renderer.calls[0]
assert call["all_messages"] == request.messages
assert call["prompt_extras"] == {
"mm_processor_kwargs": {"max_pixels": 1},
"cache_salt": "salt",
}
chat_template_kwargs = call["chat_params"].chat_template_kwargs
assert chat_template_kwargs["instruction"] == "Represent the query: "
assert chat_template_kwargs["add_generation_prompt"] is True
assert chat_template_kwargs["continue_final_message"] is False
assert "tools" not in chat_template_kwargs
assert chat_template_kwargs["tokenize"] is False
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref
from types import SimpleNamespace
import pytest
import torch
@@ -9,7 +10,10 @@ import torch
from tests.models.utils import softmax
from vllm import LLM, PoolingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.entrypoints.pooling.scoring.io_processor import CrossEncoderIOProcessor
from vllm.entrypoints.pooling.scoring.typing import ScoringData
from vllm.platforms import current_platform
from vllm.renderers import TokenizeParams
MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
PROMPT = "The chef prepared a delicious meal."
@@ -141,6 +145,45 @@ def test_max_tokens_per_doc(llm: LLM):
assert with_limit_tokens < no_limit_tokens
def test_token_type_ids_follow_post_tokenization():
processor = object.__new__(CrossEncoderIOProcessor)
processor.tokenizer = SimpleNamespace(truncation_side="right", pad_token_id=-1)
processor.renderer = SimpleNamespace(process_for_engine=lambda prompt, _: prompt)
processor.model_config = None
processor.get_score_prompt = lambda **_: (
"",
{
"prompt_token_ids": list(range(32)),
"token_type_ids": [0] * 16 + [1] * 16,
},
)
engine_inputs, pooling_params = processor._pre_process(
ScoringData(data_1=["query"], data_2=["document"]),
TokenizeParams(
max_total_tokens=None,
truncate_prompt_tokens=16,
truncation_side="left",
),
PoolingParams(task="classify", extra_kwargs={"cache_salt": "salt"}),
)
assert engine_inputs[0]["prompt_token_ids"] == list(range(16, 32))
assert pooling_params[0].extra_kwargs == {
"cache_salt": "salt",
"compressed_token_type_ids": 0,
}
engine_inputs, pooling_params = processor._pre_process(
ScoringData(data_1=["query"], data_2=["document"]),
TokenizeParams(max_total_tokens=None, pad_prompt_tokens=40),
PoolingParams(task="classify"),
)
assert engine_inputs[0]["prompt_token_ids"] == list(range(32)) + [-1] * 8
assert pooling_params[0].extra_kwargs == {"compressed_token_type_ids": 16}
def test_pooling_params(llm: LLM):
def get_outputs(use_activation):
outputs = llm.score(
@@ -161,9 +161,7 @@ def _make_pooling_serving(lora_name: str) -> _ConcretePoolingServing:
return serving
def _make_pooling_ctx(
model_name: str, serving: PoolingBaseServing
) -> PoolingServeContext:
def _make_pooling_ctx(model_name: str) -> PoolingServeContext:
mock_request = MagicMock()
mock_request.model = model_name
return PoolingServeContext(
@@ -171,9 +169,6 @@ def _make_pooling_ctx(
model_name=MODEL_NAME,
request_id="test-id",
pooling_params=PoolingParams(),
lora_request=serving._maybe_get_adapters(mock_request),
priorities=0,
prompt_extras=None,
)
@@ -181,7 +176,9 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
"""LoRA adapter name must populate ctx.lora_request without raising."""
lora_name = "bot-embed-lora"
serving = _make_pooling_serving(lora_name)
ctx = _make_pooling_ctx(lora_name, serving)
ctx = _make_pooling_ctx(lora_name)
ctx.lora_request = serving._maybe_get_adapters(ctx.request)
assert ctx.lora_request is not None
assert ctx.lora_request.lora_name == lora_name
@@ -190,6 +187,7 @@ def test_pooling_maybe_get_adapters_lora_name_sets_lora_request():
def test_pooling_maybe_get_adapters_unknown_model_raises():
"""An unrecognised model name must still raise VLLMNotFoundError."""
serving = _make_pooling_serving("some-lora")
ctx = _make_pooling_ctx("unknown-model")
with pytest.raises(VLLMNotFoundError):
_make_pooling_ctx("unknown-model", serving)
serving._maybe_get_adapters(ctx.request)
@@ -6,7 +6,6 @@ max_concurrency: 100
server_args: >-
--enforce-eager
--max-model-len 4096
--max-num-batched-tokens 32768
--safetensors-load-strategy prefetch
--moe-backend flashinfer_cutlass
--prefill-context-parallel-size 4
@@ -6,7 +6,6 @@ max_concurrency: 100
server_args: >-
--enforce-eager
--max-model-len 4096
--max-num-batched-tokens 32768
--safetensors-load-strategy prefetch
--moe-backend flashinfer_cutlass
--tensor-parallel-size 2
@@ -1,9 +0,0 @@
model_name: "poolside/Laguna-XS.2-NVFP4"
accuracy_threshold: 0.86
num_questions: 1319
num_fewshot: 5
startup_max_wait_seconds: 1200
server_args: >-
--enforce-eager
--max-model-len 4096
--trust-remote-code
@@ -2,7 +2,4 @@ model_name: "google/gemma-4-E4B-it-qat-mobile-ct"
accuracy_threshold: 0.50
num_questions: 1319
num_fewshot: 5
server_args: >-
--enforce-eager
--max-model-len 4096
--speculative-config '{"method":"mtp","model":"google/gemma-4-E4B-it-assistant","num_speculative_tokens":4}'
server_args: "--enforce-eager --max-model-len 4096"
@@ -3,4 +3,3 @@ Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
Qwen1.5-MoE-W4A16-CT.yaml
DeepSeek-V2-Lite-Instruct-FP8.yaml
Qwen3-30B-A3B-NVFP4.yaml
Laguna-XS.2-NVFP4.yaml
@@ -546,11 +546,8 @@ def test_flash_attn_accepts_handled_fp8_variants(
):
"""FlashAttentionBackend must accept the two fp8 dtypes it can actually
handle: 'fp8' (alias for fp8_e4m3fn) and 'fp8_e4m3'."""
import vllm.v1.attention.backends.fa_utils as fa_utils_mod
import vllm.v1.attention.backends.flash_attn as fa_mod
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
# The fp8 decision is made in fa_utils, using its own current_platform
# binding, so patch is_xpu there (not on flash_attn's) to stay robust to
# import order across earlier tests that patch vllm.platforms.current_platform.
monkeypatch.setattr(fa_utils_mod.current_platform, "is_xpu", lambda: True)
monkeypatch.setattr(fa_mod.current_platform, "is_xpu", lambda: True)
assert FlashAttentionBackend.supports_kv_cache_dtype(kv_cache_dtype)
@@ -11,9 +11,6 @@ from vllm._custom_ops import (
scaled_fp8_quant,
)
from vllm.platforms import current_platform
from vllm.v1.attention.ops.triton_merge_attn_states import (
mask_empty_context,
)
from vllm.v1.attention.ops.triton_merge_attn_states import (
merge_attn_states as merge_attn_states_triton,
)
@@ -76,59 +73,6 @@ DTYPES = [torch.float32, torch.half, torch.bfloat16]
all_case_info: list[tuple] = []
def test_mask_empty_context() -> None:
query_lens = torch.tensor([2] + [1] * 31 + [131, 1], dtype=torch.int32)
query_start_loc = torch.cat(
(torch.zeros(1, dtype=torch.int32), query_lens.cumsum(0))
).cuda()
context_lens = torch.tensor([4] * 32 + [0, 3], dtype=torch.int32)
context_start_loc = torch.cat(
(torch.zeros(1, dtype=torch.int32), context_lens.cumsum(0))
).cuda()
num_heads, num_tokens, head_dim = 4, 165, 16
lse = torch.randn(num_heads, num_tokens, device="cuda")
output = torch.randn(num_tokens, num_heads, head_dim, device="cuda")
# Empty-context rows carry undefined (possibly non-finite) attention output.
output[33:164] = float("nan")
expected_lse = lse.clone()
expected_lse[:, 33:164] = float("-inf")
expected_output = output.clone()
expected_output[33:164] = 0.0
mask_empty_context(lse, output, query_start_loc, context_start_loc)
torch.testing.assert_close(lse, expected_lse)
torch.testing.assert_close(output, expected_output)
@pytest.mark.parametrize("merge_fn", [merge_attn_states_cuda, merge_attn_states_triton])
@pytest.mark.parametrize("output_dtype", [torch.float32, torch.half, torch.bfloat16])
def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None:
"""When a token is empty on both sides (both LSE -inf), the 0/0 softmax
scales must not surface as NaN in the merged output."""
num_tokens, num_heads, head_size = 6, 8, 128
prefix_output = torch.zeros(
num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
)
prefix_lse = torch.randn(num_heads, num_tokens, device="cuda")
suffix_output = torch.zeros(
num_tokens, num_heads, head_size, device="cuda", dtype=output_dtype
)
suffix_lse = torch.randn(num_heads, num_tokens, device="cuda")
# Tokens 2 and 3 are empty on both sides (mask_empty_context already zeroed
# their outputs and set both LSEs to -inf).
empty = slice(2, 4)
prefix_lse[:, empty] = float("-inf")
suffix_lse[:, empty] = float("-inf")
output = torch.empty_like(prefix_output)
merge_fn(output, prefix_output, prefix_lse, suffix_output, suffix_lse)
assert not output.isnan().any()
def generate_markdown_table():
global all_case_info
table_header = (
@@ -40,11 +40,10 @@ BATCH_SIZE = 128
CONTEXT_LEN = 8192
PAGE_SIZE = 1
# On the fp8 KV-cache path the builder forwards the platform fp8 dtype (aiter
# dtypes.fp8) for both q and kv. Mirror it via current_platform.fp8_dtype()
# instead of hardcoding a literal (see #47276).
EXPECTED_Q_DTYPE = current_platform.fp8_dtype()
EXPECTED_KV_DTYPE = current_platform.fp8_dtype()
# Expected dtypes for this fold path: bf16 model dtype -> bf16 query; fp8
# KV-cache -> fp8_e4m3 kv.
EXPECTED_Q_DTYPE = torch.bfloat16
EXPECTED_KV_DTYPE = torch.float8_e4m3fn
# The split/reduce content tensors filled by get_mla_metadata_v1. work_meta_data
# is excluded: it holds raw device pointers, never equal across allocations.
-15
View File
@@ -1,15 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch
@pytest.fixture(autouse=True)
def reset_default_torch_device():
"""Several kernel tests call torch.set_default_device without restoring
it, which poisons subsequent tests in the same pytest run (e.g. CPU
tensors silently created on CUDA). Restore the factory default after
every test.
"""
yield
torch.set_default_device(None)
+19 -14
View File
@@ -33,10 +33,12 @@ from vllm.third_party.flash_linear_attention.ops.index import ( # noqa: E402
@pytest.mark.parametrize("num_seqs", [1, 5, 257])
@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32])
def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
rng_cpu = torch.Generator("cpu").manual_seed(1234)
rng = torch.Generator("cuda").manual_seed(2345)
seq_lens = torch.randint(1, 130, (num_seqs,), dtype=torch.int32, generator=rng_cpu)
seq_lens = torch.randint(
1,
130,
(num_seqs,),
dtype=torch.int32,
)
cu_seqlens = torch.zeros(num_seqs + 1, device="cuda", dtype=torch.int32)
cu_seqlens[1:] = seq_lens.to(device="cuda").cumsum(0)
total_tokens = int(cu_seqlens[-1].item())
@@ -54,9 +56,8 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
head_k_dim,
device="cuda",
dtype=dtype,
generator=rng,
)
k = torch.randn_like(q, generator=rng)
k = torch.randn_like(q)
v = torch.randn(
1,
total_tokens,
@@ -64,24 +65,29 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
head_v_dim,
device="cuda",
dtype=dtype,
generator=rng,
)
q = F.normalize(q.float(), p=2, dim=-1).to(dtype)
k = F.normalize(k.float(), p=2, dim=-1).to(dtype)
a = torch.randn(
1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng
1,
total_tokens,
num_v_heads,
device="cuda",
dtype=dtype,
)
b = torch.randn(
1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng
1,
total_tokens,
num_v_heads,
device="cuda",
dtype=dtype,
)
# Match upstream FLA GatedDeltaNet synthetic initialization:
# https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gated_deltanet.py
A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(
0, 16, generator=rng
)
A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(0, 16)
A_log = torch.log(A)
dt = torch.exp(
torch.rand(num_v_heads, device="cuda", dtype=torch.float32, generator=rng)
torch.rand(num_v_heads, device="cuda", dtype=torch.float32)
* (math.log(0.1) - math.log(0.001))
+ math.log(0.001)
)
@@ -99,7 +105,6 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
head_k_dim,
device="cuda",
dtype=state_dtype,
generator=rng,
)
* 0.05
)
@@ -6,7 +6,6 @@ import pytest
import torch
import torch.nn.functional as F
import vllm._custom_ops as ops
from vllm.model_executor.layers.fused_moe.config import (
RoutingMethodType,
get_routing_method_type,
@@ -232,119 +231,3 @@ def test_dsv4_fast_topk(
atol=2e-5,
rtol=2e-5,
)
@pytest.mark.skipif(
not current_platform.is_cuda_alike(),
reason="This test is skipped on non-CUDA platform.",
)
@pytest.mark.parametrize("use_hash", [False, True])
@pytest.mark.parametrize("use_bias", [False, True])
@pytest.mark.parametrize("use_padding_mask", [False, True])
@pytest.mark.parametrize("pad_with_nan", [False, True])
@pytest.mark.parametrize("num_experts", [128, 256, 384])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
def test_fused_topk_softplus_sqrt_padding(
use_hash: bool,
use_bias: bool,
use_padding_mask: bool,
pad_with_nan: bool,
num_experts: int,
dtype: torch.dtype,
):
"""Verify explicit padding and NaN-padded rows do not affect real rows."""
torch.manual_seed(0)
num_tokens = 8
topk = 6
indices_dtype = torch.int32
gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
padding_rows = torch.zeros(num_tokens, dtype=torch.bool, device="cuda")
padding_rows[1::2] = True
if pad_with_nan:
gating_output[padding_rows] = float("nan")
is_padding = padding_rows if use_padding_mask else None
# A negative correction bias makes explicit pad rows look selectable unless
# the kernel uses the is_padding guard.
e_score_correction_bias = None
if use_bias:
e_score_correction_bias = (
-torch.rand((num_experts,), dtype=torch.float32, device="cuda") - 1.0
)
input_ids = None
hash_indices_table = None
if use_hash:
vocab_size = 64
hash_indices_table = torch.stack(
[torch.randperm(num_experts)[:topk] for _ in range(vocab_size)]
).to(device="cuda", dtype=indices_dtype)
input_ids = torch.randint(
0, vocab_size, (num_tokens,), dtype=indices_dtype, device="cuda"
)
topk_weights = torch.empty(num_tokens, topk, dtype=torch.float32, device="cuda")
topk_ids = torch.empty(num_tokens, topk, dtype=indices_dtype, device="cuda")
token_expert_indices = torch.empty(
num_tokens, topk, dtype=torch.int32, device="cuda"
)
ops.topk_hash_softplus_sqrt(
topk_weights,
topk_ids,
token_expert_indices,
gating_output,
renormalize=True,
routed_scaling_factor=1.0,
e_score_correction_bias=e_score_correction_bias,
input_tokens=input_ids,
hash_indices_table=hash_indices_table,
is_padding=is_padding,
)
if use_padding_mask:
pad_ids = topk_ids[padding_rows]
pad_weights = topk_weights[padding_rows]
assert torch.equal(pad_ids, torch.full_like(pad_ids, -1)), (
f"Explicit pad rows should contain only -1 ids, got {pad_ids.tolist()}"
)
assert (pad_weights == 0).all(), (
"Explicit pad rows should have all-zero weights, "
f"got {pad_weights.tolist()}"
)
if pad_with_nan:
nan_pad_weights = topk_weights[padding_rows]
assert torch.isfinite(nan_pad_weights).all(), (
f"NaN-padded rows have non-finite weights, got {nan_pad_weights.tolist()}"
)
assert (nan_pad_weights == 0).all(), (
"NaN-padded rows should have all-zero weights, "
f"got {nan_pad_weights.tolist()}"
)
topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
gating_output=gating_output,
topk=topk,
renormalize=True,
routed_scaling_factor=1.0,
e_score_correction_bias=e_score_correction_bias,
input_ids=input_ids,
hash_indices_table=hash_indices_table,
)
rows_to_compare = torch.ones(num_tokens, dtype=torch.bool, device="cuda")
if use_padding_mask or pad_with_nan:
rows_to_compare = ~padding_rows
sorted_ref_ids, idx_ref = topk_ids_ref[rows_to_compare].sort(dim=-1)
sorted_ids, idx_ops = topk_ids[rows_to_compare].sort(dim=-1)
torch.testing.assert_close(
sorted_ref_ids, sorted_ids.to(sorted_ref_ids.dtype), atol=0, rtol=0
)
sorted_w_ref = topk_weights_ref[rows_to_compare].gather(1, idx_ref)
sorted_w = topk_weights[rows_to_compare].gather(1, idx_ops)
torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2)
+3 -5
View File
@@ -21,9 +21,9 @@ def test_gather_cache_oob():
seq_starts causes the block_table offset to read out of bounds.
"""
batch_size = 1
block_size = 64
# The kernel only supports the MLA entry sizes.
entry_size = 576
entry_size = 128
block_table = torch.tensor([[1, 2]], dtype=torch.int32, device="cuda")
@@ -34,7 +34,6 @@ def test_gather_cache_oob():
seq_len = 65
cu_seq_lens = torch.tensor([0, seq_len], dtype=torch.int32, device="cuda")
token_to_seq = torch.zeros(seq_len, dtype=torch.int32, device="cuda")
# src_cache: [num_blocks, block_size, entry_size]
num_blocks = 5
@@ -52,8 +51,7 @@ def test_gather_cache_oob():
dst,
block_table,
cu_seq_lens,
token_to_seq,
seq_len,
batch_size,
"auto", # kv_cache_dtype
scale,
seq_starts,
-21
View File
@@ -13,7 +13,6 @@ These tests cover:
"""
import math
from types import SimpleNamespace
import pytest
import torch
@@ -28,7 +27,6 @@ from vllm.models.deepseek_v4.common.ops.fused_compress_quant_cache import (
_fused_kv_compress_norm_rope_insert_indexer_mxfp4_attn,
_launch_two_stage_sparse_attn_compressor,
)
from vllm.models.deepseek_v4.compressor import _get_c128_boundary
from vllm.platforms import current_platform
from .test_fused_indexer_q_rope_quant import quantize_to_mxfp4
@@ -60,25 +58,6 @@ def _ue8m0_reference(x: torch.Tensor, block_size: int, fp8_max: float):
return x_fp8, scales
@pytest.mark.parametrize(
("starts", "query_start_loc", "expected"),
[
([0], [0, 127], False),
([0], [0, 128], True),
([127], [0, 1], True),
([128], [0, 127], False),
([1, 255], [0, 1, 2], True),
(None, [0, 1], None),
],
)
def test_get_c128_boundary(starts, query_start_loc, expected):
metadata = SimpleNamespace(
_num_computed_tokens_cpu=None if starts is None else torch.tensor(starts),
query_start_loc_cpu=torch.tensor(query_start_loc),
)
assert _get_c128_boundary(metadata) is expected
# ── Test A: DeepseekV4 Attention path ──────────────────────────────────────────────
+1 -1
View File
@@ -264,7 +264,7 @@ def test_block_mask_direct_vs_slow_path():
device = torch.device("cuda")
vllm_config = create_vllm_config(
model_name="Qwen/Qwen2.5-1.5B-Instruct", block_size=16, max_model_len=1024
model_name="meta-llama/Meta-Llama-3-8B", block_size=16, max_model_len=1024
)
kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
+4 -13
View File
@@ -726,19 +726,14 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
This catches stride/layout bugs that only manifest when the einsum
kernel actually consumes the quantized activations.
"""
from deep_gemm.utils.math import ceil_div
from vllm.utils.deep_gemm import (
fp8_einsum,
is_deep_gemm_supported,
per_block_cast_to_fp8,
transform_sf_into_required_layout,
)
if not is_deep_gemm_supported():
pytest.skip("DeepGEMM not supported on this platform")
def ceil_div(a: int, b: int) -> int:
return (a + b - 1) // b
heads_per_group = num_heads // n_groups
d = heads_per_group * HEAD_DIM
o_lora_rank = 1024
@@ -814,12 +809,8 @@ def test_einsum_end_to_end(num_tokens, num_heads, n_groups):
# -- Checks --
# Einsum output: Triton and CUDA both rotate in fp32 now, so diffs
# come from fp32 ordering and UE8M0 boundary shifts only.
# Use relative diff (same metric as deep_gemm.testing.calc_diff).
def calc_diff(x, y):
x, y = x.double(), y.double()
denominator = (x * x + y * y).sum()
sim = 2 * (x * y).sum() / denominator
return 1 - sim
# Use relative diff (same metric as test_fp8_einsum.py).
from deep_gemm.testing import calc_diff
z_diff = calc_diff(z_fused, z_ref)
assert z_diff < 0.01, (
@@ -84,22 +84,6 @@ def norm_rope_ref(x, weight, positions, cos_sin_cache, eps):
return roped
def assert_fp8_cache_close(kv_cache, expected_kv_cache):
"""Compare two e4m3 caches allowing 1 ulp.
On CUDA the fused kernel quantizes K from its fp32 intermediate, while the
reshape_and_cache_flash reference quantizes the bf16-materialized value, so
rounding-boundary values may differ by one e4m3 code.
"""
byte_diff = (kv_cache.int() - expected_kv_cache.int()).abs()
got = kv_cache.view(torch.float8_e4m3fn).float()
exp = expected_kv_cache.view(torch.float8_e4m3fn).float()
ok = (byte_diff <= 1) | ((got == 0) & (exp == 0))
assert bool(ok.all()), (
f"fp8 cache differs by more than 1 ulp in {int((~ok).sum())} elements"
)
# ── Test 1: dense mode (norm+rope only, no index, no insert) ─────────────────
@@ -281,7 +265,7 @@ def test_sparse_full(num_tokens, block_size, kv_cache_dtype):
scale,
scale,
)
assert_fp8_cache_close(kv_cache, expected_kv_cache)
torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
else:
for t in range(num_tokens):
s = slot_mapping[t].item()
@@ -399,7 +383,7 @@ def test_sparse_skip_index_branch(num_tokens, block_size, kv_cache_dtype):
scale,
scale,
)
assert_fp8_cache_close(kv_cache, expected_kv_cache)
torch.testing.assert_close(kv_cache, expected_kv_cache, rtol=0, atol=0)
else:
k_ref_h = k_ref.view(num_tokens, num_kv_heads, HEAD_DIM)
v_ref_h = v_in.view(num_tokens, num_kv_heads, HEAD_DIM)
@@ -41,12 +41,11 @@ def test_fused_recurrent_packed_decode_matches_reference(
A_log = torch.randn((HV,), device=device, dtype=dtype)
dt_bias = torch.randn((HV,), device=device, dtype=dtype)
# Continuous batching indices (include PAD_SLOT_ID=-1 cases). Index 0 is
# reserved as NULL_BLOCK_ID (CUDA graph padding), so valid slots start at 1.
ssm_state_indices = torch.arange(1, B + 1, device=device, dtype=torch.int32)
# Continuous batching indices (include PAD_SLOT_ID=-1 cases).
ssm_state_indices = torch.arange(B, device=device, dtype=torch.int32)
ssm_state_indices[-3:] = -1
state0 = torch.randn((B + 1, HV, V, K), device=device, dtype=dtype)
state0 = torch.randn((B, HV, V, K), device=device, dtype=dtype)
state_ref = state0.clone()
state_packed = state0.clone()
@@ -95,8 +94,5 @@ def test_fused_recurrent_packed_decode_matches_reference(
atol = 2e-2 if dtype != torch.float32 else 1e-4
rtol = 1e-2 if dtype != torch.float32 else 1e-4
# Output rows for PAD_SLOT_ID entries are never written (uninitialized in
# both paths), so compare only the valid rows.
valid = ssm_state_indices > 0
torch.testing.assert_close(out_packed[valid], out_ref[valid], rtol=rtol, atol=atol)
torch.testing.assert_close(out_packed, out_ref, rtol=rtol, atol=atol)
torch.testing.assert_close(state_packed, state_ref, rtol=rtol, atol=atol)
@@ -58,12 +58,10 @@ def test_fused_sigmoid_gating_delta_rule_update_non_spec(
dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
# Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
# state indices start at 1.
ssm_state = torch.rand(
total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
)
state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[:num_tokens]
state_indices = torch.randperm(total_entries, dtype=torch.int32)[:num_tokens]
cu_seqlens = torch.arange(0, num_tokens + 1, dtype=torch.int32)
beta = b.sigmoid()
@@ -146,14 +144,13 @@ def test_fused_sigmoid_gating_delta_rule_update_spec(
dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
# Entry 0 is reserved as NULL_BLOCK_ID (CUDA graph padding), so valid
# state indices start at 1.
ssm_state = torch.rand(
total_entries + 1, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
)
state_indices = (torch.randperm(total_entries, dtype=torch.int32) + 1)[
:num_tokens
].view(num_reqs, num_speculative_tokens + 1)
state_indices = torch.randperm(
total_entries,
dtype=torch.int32,
)[:num_tokens].view(num_reqs, num_speculative_tokens + 1)
num_accepted_tokens = torch.randint(
1, num_speculative_tokens + 1, (num_reqs,), dtype=torch.int32
)

Some files were not shown because too many files have changed in this diff Show More